feat(webgl): update pool2d

Author: BenAnn

packages/paddlejs-backend-webgl/src/ops/shader/pool2d_max.ts

@@ -8,12 +8,22 @@ function recoverShape({ total_shape, channel, height_shape, width_shape }) {
 
 function mainFunc(
     { origin },
-    { strides = [], paddings = [], ksize, global_pooling }
+    { strides = [], paddings = [], ksize, global_pooling, runtime }
 ) {
     const [stride_v = 1, stride_h = 1] = strides;
     const [padTop = 0, padLeft = 0] = paddings;
     const [ksize_x, ksize_y] = ksize;
     const originShape = recoverShape(origin);
+    let computedIndex = '';
+    let outputCode = 'setOutput(float(res));';
+    if (runtime === 0 && global_pooling === true) {
+        computedIndex = `
+            if (curr > res) {
+                index = ${originShape[2] * originShape[3]} * out_pos[1] + ${originShape[3]} * oy + ox;
+            }
+        `;
+        outputCode = 'setOutput(float(index));';
+    }
 
     return `
     // start函数
@@ -47,20 +57,11 @@ function mainFunc(
                 }
                 // origin数据
                 float curr = getValueFromTensorPos_origin(out_pos[0], out_pos[1], oy, ox);
-                if (layer_run_time == 0 && ${global_pooling === true}) {
-                    if (curr > res) {
-                        index = ${originShape[2] * originShape[3]} * out_pos[1] + ${originShape[3]} * oy + ox;
-                    }
-                }
+                ${computedIndex}
                 res = max(res, curr);
             }
         }
-        if (layer_run_time == 0 && ${global_pooling === true}) {
-            setOutput(float(index));
-        }
-        else {
-            setOutput(float(res));
-        }
+        ${outputCode}
     }
     `;
 }
@@ -70,7 +71,8 @@ export default {
         'strides',
         'paddings',
         'ksize',
-        'global_pooling'
+        'global_pooling',
+        'runtime'
     ],
     textureFuncConf: {
         origin: ['getValueFromTensorPos']

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

@@ -21,7 +21,7 @@ const inputParams = [
     'limit',
     'channel',
     'total_shape',
-    'numbers_shape',
+    'numbers_shape'
 ];
 
 const outParams = [
@@ -44,7 +44,7 @@ const baseParams = {
     ]
 };
 
-function getTensorParams(inputTensors: Tensor[], ownParams: [], fShaderParams: object): opInfo {
+function getTensorParams(inputTensors: Tensor[], ownParams: [], fShaderParams: object, runtime: number): opInfo {
     const tensorsParams = {};
     const opParams = {};
     const tensorNames = [] as string[];
@@ -97,6 +97,9 @@ function getTensorParams(inputTensors: Tensor[], ownParams: [], fShaderParams: o
     if (fShaderParams['active_function']) {
         opParams['active_function'] = fShaderParams['active_function'];
     }
+
+    opParams['runtime'] = runtime;
+
     return { textureParams: tensorsParams, opParams, active_function: fShaderParams['active_function'] };
 }
 

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

@@ -20,7 +20,9 @@ export default function buildShader(textureConf, type, inputTensors, fShaderPara
         const { params = {}, mainFunc, textureFuncConf = {}, commonFuncConf } = ops[opName];
 
         // textureList: [filter, origin, bias]
-        const { textureParams, opParams, active_function } = getTensorParams(inputTensors, params, fShaderParams);
+        const { textureParams, opParams, active_function } = getTensorParams(
+            inputTensors, params, fShaderParams, runtime
+        );
 
         const prefixCode = genPrefixCode(textureConf);