Merge pull request PaddlePaddle#166 from JingyuanZhang/master

 feat(core): add webgl_pack_out feature

packages/paddlejs-backend-webgl/src/backend.ts

@@ -4,13 +4,13 @@
  */
 
 import { PaddlejsBackend, env } from '@paddlejs/paddlejs-core';
-import { OpData, Query } from './types';
+import { ModelVar, OpData, Query } from './types';
 import { GLHelper, EShaderType } from './webgl/WebGLUtils';
 import { GLTexture, TextureConfig } from './webgl/WebGLTexture';
 import { vShaderSource, vShaderData } from './ops/vShader';
 import buildShader from './webgl/buildShader';
 import GLProgram from './webgl/WebGLProgram';
-import { nhwc2nchw } from './utils/dataProcess';
+import { getSizeFromShape, nhwc2nchw } from './utils/dataProcess';
 import queryProcess from './utils/queryProcess';
 
 
@@ -135,14 +135,20 @@ export default class WebGLBackend extends PaddlejsBackend {
         }
     }
 
-    async read(): Promise<number[]> {
+    async read(fetchInfo: ModelVar): Promise<number[]> {
         const pbo = this.createPBO();
         await this.createAndWaitForFence();
         const result = this.downloadFoat32TensorFromBuffer(pbo);
-        let shape = (this.program as GLProgram).shape as number[];
-        if (env.get('debug') && env.get('shape')) {
-            shape = env.get('shape');
+
+        let shape = fetchInfo ? fetchInfo.shape : [];
+        if (env.get('webgl_pack_output')) {
+            return result.slice(0, getSizeFromShape(shape));
         }
+
+        shape = env.get('debug') && env.get('shape')
+            ? env.get('shape')
+            : (this.program as GLProgram).shape;
+
         const [N, C, H, W] = shape;
         const nhwcFetchShape = [N, H, W, C];
         return nhwc2nchw(result, nhwcFetchShape);
@@ -224,7 +230,11 @@ export default class WebGLBackend extends PaddlejsBackend {
             gl2.bindBuffer(gl2.PIXEL_PACK_BUFFER, buffer);
             gl2.getBufferSubData(gl2.PIXEL_PACK_BUFFER, 0, pixels);
             gl2.bindBuffer(gl2.PIXEL_PACK_BUFFER, null);
+
             const result: number[] = [];
+            if (env.get('webgl_pack_output')) {
+                return Array.from(pixels);
+            }
             for (let i = 0; i < this.width_texture_out * this.height_texture_out; i++) {
                 result.push(pixels[4 * i]);
             }
@@ -285,7 +295,14 @@ export default class WebGLBackend extends PaddlejsBackend {
         return this.frameBuffer;
     }
 
-    render(data: any = [], iLayer: number = 0, isRendered: Boolean = false, index: number, isPacked: Boolean = false, modelName: string) {
+    render(
+        data: any = [],
+        iLayer: number = 0,
+        isRendered: Boolean = false,
+        index: number,
+        isPacked: Boolean = false,
+        modelName: string
+    ) {
         const gl = this.gl;
         const that = this;
         let textureIndex = 0;

packages/paddlejs-backend-webgl/src/ops/index.ts

@@ -27,8 +27,6 @@ import bilinear_interp_v2 from './shader/bilinear_interp_v2';
 import transpose2 from './shader/transpose2';
 import softmax from './shader/softmax';
 import dynamic from './shader/dynamic';
-import unpacked_2_packed from './shader/unpacked_2_packed';
-import packed_2_unpacked from './shader/packed_2_unpacked';
 import elementwise_mul from './shader/elementwise_mul';
 import elementwise_div from './shader/elementwise_div';
 import arg_max from './shader/arg_max';
@@ -50,6 +48,10 @@ import elementwise_pow from './shader/elementwise_pow';
 import elementwise_sub from './shader/elementwise_sub';
 import cast from './shader/cast';
 
+import { pack_out, nhwc_2_nchw, unpacked_2_packed, packed_2_unpacked } from './shader/custom';
+
+
+
 const ops = {
     arg_max,
     arg_min,
@@ -105,7 +107,9 @@ const ops = {
     squeeze2,
     pad3d,
     bilinear_interp_v2,
-    shuffle_channel
+    shuffle_channel,
+    pack_out,
+    nhwc_2_nchw
 };
 export {
     ops

packages/paddlejs-backend-webgl/src/ops/shader/custom/index.ts

@@ -0,0 +1,11 @@
+import nhwc_2_nchw from './nhwc_2_nchw';
+import pack_out from './pack_out';
+import unpacked_2_packed from './unpacked_2_packed';
+import packed_2_unpacked from './packed_2_unpacked';
+
+export {
+    nhwc_2_nchw,
+    pack_out,
+    unpacked_2_packed,
+    packed_2_unpacked
+};

packages/paddlejs-backend-webgl/src/ops/shader/custom/nhwc_2_nchw.ts

@@ -0,0 +1,36 @@
+/**
+ * @file fetch
+ */
+
+function mainFunc(
+    { origin, out },
+    {}
+) {
+    return `
+    void main() {
+        ivec4 oPos = getOutputTensorPos();
+        // 输出坐标转换为输入坐标
+        int sumVal = oPos.a * ${out.channel}
+            + oPos.b * ${out.width_shape} * ${out.channel}
+            + oPos.g
+            + oPos.r * ${out.channel} * ${out.width_shape} * ${out.height_shape};
+        ivec4 new_oPos = transferFromNHWCtoNCHW(
+            sumVal,
+            ${origin.channel},
+            ${origin.width_shape},
+            ${origin.height_shape},
+            ${origin.total_shape}
+        );
+        float o = getValueFromTensorPos_origin(new_oPos.r, new_oPos.g, new_oPos.b, new_oPos.a);
+        setOutput(float(o));
+    }
+    `;
+}
+export default {
+    mainFunc,
+    params: [],
+    textureFuncConf: {
+        origin: ['getValueFromTensorPos']
+    },
+    commonFuncConf: ['transferFromNHWCtoNCHW']
+};

packages/paddlejs-backend-webgl/src/ops/shader/custom/pack_out.ts

@@ -0,0 +1,61 @@
+/**
+ * @file pack out
+ */
+
+function mainFunc(
+    { origin },
+    {}
+) {
+    const width_texture_origin = origin.width_texture;
+    const height_texture_origin = origin.height_texture;
+    return `
+
+    vec2 getOriginCoord(float x, float y) {
+        if (x > float(${width_texture_origin})) {
+            int num = int(x / float(${width_texture_origin}));
+            x = mod(x, float(${width_texture_origin}));
+            y = y + float(num);
+        }
+        return vec2(x, y);
+    }
+
+    float getClipedCoordRed(vec2 xy) {
+        return TEXTURE2D(
+            texture_origin,
+            vec2(float(xy.x / float(${width_texture_origin})), float(xy.y / float(${height_texture_origin})))
+        ).r;
+    }
+
+    // start函数
+    void main() {
+        
+        vec2 outCoord = vCoord.xy * _2d_shape_texture_out;
+        vec4 out4;
+        float x = floor(outCoord.x) * 4.0;
+        float y = floor(outCoord.y) * 4.0 + 0.5;
+        float x0 = x + 0.5;
+        float x1 = x + 1.5;
+        float x2 = x + 2.5;
+        float x3 = x + 3.5;
+
+        vec2 xy0 = getOriginCoord(x0, y);
+        vec2 xy1 = getOriginCoord(x1, y);
+        vec2 xy2 = getOriginCoord(x2, y);
+        vec2 xy3 = getOriginCoord(x3, y);
+
+        float r = getClipedCoordRed(xy0);
+        float g = getClipedCoordRed(xy1);
+        float b = getClipedCoordRed(xy2);
+        float a = getClipedCoordRed(xy3);
+
+        setPackedOutput(vec4(r, g, b, a));
+    }
+    `;
+}
+export default {
+    mainFunc,
+    params: [],
+    textureFuncConf: {
+        origin: ['getValueFromTensorPosPacking', 'getValueFromTensorPos']
+    }
+};

packages/paddlejs-backend-webgl/src/types.ts

@@ -1,5 +1,5 @@
 import type Tensor from '@paddlejs/paddlejs-core/src/opFactory/tensor';
-import { OpData } from '@paddlejs/paddlejs-core/src/commons/interface';
+import { OpData, ModelVar } from '@paddlejs/paddlejs-core/src/commons/interface';
 
 interface Query {
     name: string;
@@ -10,5 +10,6 @@ interface Query {
 export {
     Tensor,
     OpData,
-    Query
+    Query,
+    ModelVar
 };

packages/paddlejs-backend-webgl/src/utils/dataProcess.ts

@@ -35,7 +35,12 @@ function reduceShape(shape: number[]): number[] {
     return s;
 }
 
+function getSizeFromShape(shape: number[]): number {
+    return shape.reduce((acc, cur) => acc * cur, 1);
+}
+
 export {
+    getSizeFromShape,
     nhwc2nchw,
     reduceShape
 };

packages/paddlejs-backend-webgl/test/model/modelTest.js

@@ -13,7 +13,8 @@ async function run() {
         }
     });
 
-    env.set('debug', true);
+    env.set('webgl_pack_output', true);
+
     const preheatRes = await runner.init();
     console.log(preheatRes);
     window.weightMap = runner.weightMap;

packages/paddlejs-backend-webgl/test/op/data/pack_out.json

@@ -0,0 +1,25 @@
+{
+    "ops": [
+        {
+            "attrs": {},
+            "inputs": {
+                "X": ["concat.tmp_0"]
+            },
+            "outputs": {
+                "Out": ["concat.tmp_2"]
+            },
+            "type": "pack_out"
+        }
+    ],
+    "vars": [
+        {
+            "data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
+            "name": "concat.tmp_0",
+            "shape": [1, 1, 6, 2]
+        },
+        {
+            "name": "concat.tmp_2",
+            "shape": [1, 1, 2, 2]
+        }
+    ]
+}

packages/paddlejs-core/src/transform/index.ts

@@ -2,6 +2,7 @@ import TexturePacking from './texturePacking';
 import FormatInputsX from './formatInputsX';
 import type Transformer from './transformer';
 import SplitOp from './splitOp';
+import PackOutOp from './packOutOp';
 
 interface TransformerAction {
     preTransforms: Transformer[];
@@ -11,7 +12,8 @@ interface TransformerAction {
 
 const actions: TransformerAction = {
     preTransforms: [
-        new SplitOp()
+        new SplitOp(),
+        new PackOutOp()
     ],
     transforms: [
         new FormatInputsX(),

packages/paddlejs-core/src/transform/packOutOp.ts

@@ -0,0 +1,81 @@
+/**
+ * @file nhwc 2 nchw & pack out
+ */
+
+import env from '../env';
+import { ModelOp, ModelVar } from '../commons/interface';
+import { formatShape } from '../opFactory/utils';
+import Transformer from './transformer';
+
+const FINAL_PACK_OP_NAME = 'fetch_pack';
+const FINAL_NCHW_OP_NAME = 'final_nchw';
+
+export default class Fetch extends Transformer {
+    constructor() {
+        super('Fetch');
+    }
+
+    transform(...args: any) {
+        if (!env.get('webgl_pack_output')) {
+            return;
+        }
+        const [ops, vars] = args;
+        const fetchOp = ops.find(item => item.type === 'fetch');
+        const [inputName] = fetchOp.inputs.X;
+        const fetchInputVar = vars.find(item => item.name === inputName);
+        const [n, h, w, c] = formatShape(fetchInputVar.shape);
+
+        // transform data from nhwc to nchw
+        const nchwOp: ModelOp = {
+            attrs: {},
+            inputs: {
+                X: [inputName]
+            },
+            outputs: {
+                Y: [FINAL_NCHW_OP_NAME]
+            },
+            type: 'nhwc_2_nchw'
+        };
+
+        // pack out texture
+        const packOutOp: ModelOp = {
+            attrs: {},
+            inputs: {
+                X: [FINAL_NCHW_OP_NAME]
+            },
+            outputs: {
+                Y: [FINAL_PACK_OP_NAME]
+            },
+            type: 'pack_out'
+        };
+
+        // make nchw op var
+        const nchwVar = {
+            name: FINAL_NCHW_OP_NAME,
+            shape: [n, c, h, w],
+            persistable: false
+        };
+
+        const pack_width = c * w;
+        const pack_height = Math.ceil(n * h / 4);
+        // make pack op var
+        const packOutVar = {
+            name: FINAL_PACK_OP_NAME,
+            shape: [1, 1, pack_height, pack_width],
+            persistable: false
+        };
+
+        const changed_fetch_name = `${inputName}_fetch`;
+        fetchOp.inputs.X = [changed_fetch_name];
+        // save origin fetch op info
+        const changedFetchVar: ModelVar = {
+            name: changed_fetch_name,
+            shape: fetchInputVar.shape,
+            persistable: false
+        };
+
+        ops.push(...[nchwOp, packOutOp]);
+        vars.push(...[nchwVar, packOutVar, changedFetchVar]);
+    }
+}
+

packages/paddlejs-core/src/transform/texturePacking.ts

@@ -3,6 +3,7 @@
  */
 
 import { GLOBALS } from '../globals';
+import env from '../env';
 import OpExecutor from '../opFactory/opExecutor';
 import Transformer from './transformer';
 
@@ -109,10 +110,11 @@ export default class TexturePacking extends Transformer {
     }
 
     transform(...args: any) {
-        // wegbl backend 单独处理
-        if (GLOBALS.backend !== 'webgl') {
+        // wegbl backend 单独处理，开启 webgl pack flag 才处理
+        if (GLOBALS.backend !== 'webgl' || !env.get('webgl_pack_channel')) {
             return;
         }
+
         const [originOp, vars, opsMap] = args;
 
         if (!(packedOpConditions[originOp.type] && packedOpConditions[originOp.type](originOp, vars))) {