license import { GPGPUProgram , useShapeUniforms } from './gpgpu_math' ; export class MatMulPackedProgram implements GPGPUProgram { variableNames = [ 'matrixA' , 'matrixB' ] ; packedInputs = true ; packedOutput = true ; outputShape : number [ ] ; userCode : string ; enableShapeUniforms : boolean ; constructor ( aShape , bShape , outputShape , transposeA = false , transposeB = false , addBias = false , activation = null , hasPreluActivation = false , hasLeakyreluActivation = false ) { this . outputShape = outputShape ; this . enableShapeUniforms = useShapeUniforms ( this . outputShape . length ) ; const sharedDim = transposeA ? aShape [ 1 ] : aShape [ 2 ] ; const sharedDimensionPacked = Math . ceil ( sharedDim / 2 ) ; const aSample = transposeA ? 'i * 2, rc.y' : 'rc.y, i * 2' ; const bSample = transposeB ? 'rc.z, i * 2' : 'i * 2, rc.z' ; const aSwizzle = transposeA ? [ 'a.xxyy' , 'a.zzww' ] : [ 'a.xxzz' , 'a.yyww' ] ; const bSwizzle = transposeB ? [ 'b.xzxz' , 'b.ywyw' ] : [ 'b.xyxy' , 'b.zwzw' ] ; let activationSnippet = '' , applyActivationSnippet = '' ; if ( activation ) { if ( hasPreluActivation ) { activationSnippet = `vec4 activation(vec4 a) {
          vec4 b = getPreluActivationWeightsAtOutCoords();
          ${ activation }
        }` ; } else if ( hasLeakyreluActivation ) { activationSnippet = `vec4 activation(vec4 a) {
          vec4 b = getLeakyreluAlphaAtOutCoords();
          ${ activation }
        }` ; } else { activationSnippet = `vec4 activation(vec4 x) {
          ${ activation }
        }` ; } applyActivationSnippet = `result = activation(result);` ; } const addBiasSnippet = addBias ? 'result += getBiasAtOutCoords();' : '' ; if ( addBias ) { this . variableNames . push ( 'bias' ) ; } if ( hasPreluActivation ) { this . variableNames . push ( 'preluActivationWeights' ) ; } if ( hasLeakyreluActivation ) { this . variableNames . push ( 'leakyreluAlpha' ) ; } let batchASnippet = 'rc.x' ; let batchBSnippet = 'rc.x' ; if ( aShape [ 0 ] < bShape [ 0 ] ) { batchASnippet = `int(min(float(rc.x), ${ aShape [ 0 ] - 1 }.))` ; } else if ( bShape [ 0 ] < aShape [ 0 ] ) { batchBSnippet = `int(min(float(rc.x), ${ bShape [ 0 ] - 1 }.))` ; } this . userCode = `
      ${ activationSnippet }
      // Don't use uniform for sharedDimensionPacked for performance.
      const float sharedDimension = ${ sharedDimensionPacked }.0;

      vec4 dot2x2ARowBCol(ivec3 rc) {
        vec4 result = vec4(0);
        for (int i = 0; i < ${ sharedDimensionPacked }; i++) {
          int batchA = ${ batchASnippet };
          int batchB = ${ batchBSnippet };
          vec4 a = getMatrixA(batchA, ${ aSample });
          vec4 b = getMatrixB(batchB, ${ bSample });

          // These swizzled products need to be separately added.
          // See: https://github.com/tensorflow/tfjs/issues/1735
          result += (${ aSwizzle [ 0 ] } * ${ bSwizzle [ 0 ] });
          result += (${ aSwizzle [ 1 ] } * ${ bSwizzle [ 1 ] });
        }
        return result;
      }

      void main() {
        ivec3 rc = getOutputCoords();
        vec4 result = dot2x2ARowBCol(rc);

        ${ addBiasSnippet }

        ${ applyActivationSnippet }

        setOutput(result);
      }
    ` ; } }