fix(webgl): fix slice stack op run in webgl1.0 GLSL ES 1.0

Author: JingyuanZhang

packages/paddlejs-backend-webgl/src/ops/atom/common_utils.ts

@@ -0,0 +1,38 @@
+/**
+ * @file common utils
+ */
+
+import { env } from '@paddlejs/paddlejs-core';
+
+
+enum ArrTypeEnum {
+    INT_TYPE = 'int',
+    FLOAT_TYPE = 'float'
+}
+
+// GLSL ES 3.00 支持 arr =>  int arr = int[](x, x, x,... x);
+// GLSL ES 1.0 (1.0) 不支持  array constructor
+// '[]' : array constructor supported in GLSL ES 3.00 and above only
+const initializeGLSLArr = (arr: Array<Number>, type: ArrTypeEnum) => {
+    if (env.get('webglVersion') !== 2) {
+        return arr.reduce((acc, cur, index) => {
+            const tmp = index < arr.length - 1 ? `${cur}, ` : `${cur});`;
+            return acc + tmp;
+        }, `${type} arr[] = ${type}[](`);
+    }
+
+    const arr_value = arr.reduce((acc, cur, index) => {
+        return acc + `
+            arr[${index}] = ${cur};`;
+    }, '');
+
+    return `
+        ${type} arr[${arr.length}];
+        ${arr_value}
+    `;
+};
+
+export {
+    initializeGLSLArr,
+    ArrTypeEnum
+};

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

@@ -4,6 +4,10 @@
  * @example x = [[1,2,3,4],[5,6,7,8]] axes=[1] starts=[2] ends=[3] => out [3,7]
  */
 
+import { env } from '@paddlejs/paddlejs-core';
+import { initializeGLSLArr, ArrTypeEnum } from '../atom/common_utils';
+
+
 function mainFunc(
     {
         out, origin
@@ -70,23 +74,39 @@ function mainFunc(
             }
         }
     }
-    // 生成 glsl arr =>  int arr = int[](x, x, x,... x);
-    const glslIndexArr = res_pos.reduce((acc, cur, index) => {
-        const tmp = index < res_pos.length - 1 ? `${cur}, ` : `${cur});`;
-        return acc + tmp;
-    }, 'int arr[] = int[](');
 
+    const glslIndexArr = initializeGLSLArr(res_pos, ArrTypeEnum.INT_TYPE);
+
+    const ifConditions = res_pos.reduce((acc, _, idx) => {
+        const ifCondition = idx === 0
+            ? `
+                int index = 0;
+                if (sumVal == ${idx}) {
+                    index = arr[${idx}];
+                }`
+            : `
+                else if (sumVal == ${idx}) {
+                    index = arr[${idx}];
+                }
+            `;
+        return acc + ifCondition;
+    }, '');
+
+    const getValueFromArrIndex = env.get('webglVersion') === 2
+        ? 'int index = arr[sumVal];'
+        : ifConditions;
     return `
     void main(void) {
         ivec4 oPos = getOutputTensorPos();
         ${glslIndexArr}
+
         // 输出坐标转换为输入坐标
         int sumVal = oPos.a
             + oPos.b * ${out.width_shape}
             + oPos.g * ${out.height_shape} * ${out.width_shape}
             + oPos.r * ${out.channel} * ${out.width_shape} * ${out.height_shape};
         
-        int index = arr[sumVal];
+        ${getValueFromArrIndex}
 
         float res = 0.0;
         ivec4 co = getTensorPosFromArrayIndex_origin(index);

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

@@ -66,14 +66,13 @@ function mainFunc(
             + oPos.g * ${out.height_shape} * ${out.width_shape}
             + oPos.r * ${out.channel} * ${out.width_shape} * ${out.height_shape};
 
-        int index = sumVal % ${pre_every_num};
+        int index = calMod(sumVal, ${pre_every_num});
 
         int layer = sumVal / ${pre_every_num};
 
         int i = index / ${post};
-        int j = index % ${post} + layer * ${post};
+        int j = calMod(index, ${post}) + layer * ${post};
 
-        
         float o = 0.0;
         ${getMultiInputsValue}
         setOutput(float(o));

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

@@ -86,9 +86,11 @@
             "name": "stack_out",
             "shape": [3, 3, 3],
             "data": [
-                1, 2, 3, 4, 5, 6, 7, 8, 9,
-                11, 12, 13, 14, 15, 16, 17, 18, 19,
-                21, 22, 23, 24, 25, 26, 27, 28, 29
+                1, 11, 21, 2, 12, 22,
+                3, 13, 23, 4, 14, 24,
+                5, 15, 25, 6, 16, 26,
+                7, 17, 27, 8, 18, 28,
+                9, 19, 29
             ]
         }
     ]