license import { backend_util } from '@tensorflow/tfjs-core' ; import { GPGPUProgram } from './gpgpu_math' ; export class SegmentOpProgram implements GPGPUProgram { variableNames = [ 'x' , 'segmentIds' ] ; outputShape : number [ ] ; userCode : string ; constructor ( segOpInfo , segOpType ) { const windowSize = segOpInfo . windowSize ; const batchSize = segOpInfo . batchSize ; const inSize = segOpInfo . inSize ; const numSegments = segOpInfo . numSegments ; const outSize = numSegments * Math . ceil ( inSize / windowSize ) ; this . outputShape = [ batchSize , outSize ] ; const initializationValue = '0.0' ; const returnValue = `sumValue` ; const windowSizeNearestVec4 = Math . floor ( windowSize / 4 ) * 4 ; const windowSizeVec4Remainder = windowSize % 4 ; const updateSnippet = `
        sumValue += dot(values, segFilter);
    ` ; let checkValueOutOfBounds = '' ; if ( inSize % windowSize > 0 ) { checkValueOutOfBounds = `
        if (inIdx < 0 || inIdx >= ${ inSize }) {
          return initializationValue;
        }
      ` ; } let checkSegmentIdOutOfBounds = '' ; if ( inSize % windowSize > 0 ) { checkSegmentIdOutOfBounds = `
        if (inIdx < 0 || inIdx >= ${ inSize }) {
          return -1.0;
        }
      ` ; } this . userCode = `
      const float initializationValue = ${ initializationValue };

      float getValue(int batch, int inIdx) {
        ${ checkValueOutOfBounds }
        return getX(batch, inIdx);
      }

      float getSegmentIdAtIndex(int inIdx) {
        ${ checkSegmentIdOutOfBounds }
        return getSegmentIds(inIdx);
      }

      void main() {
        ivec2 coords = getOutputCoords();
        int batch = coords[0];
        int outIdx = coords[1];
        int inOffset = int(floor(float(outIdx) / float(
          ${ numSegments })) * float(${ windowSize }));
        int currentSeg = int(mod(float(outIdx), float(${ numSegments })));

        float sumValue = 0.0;

        for (int i = 0; i < ${ windowSizeNearestVec4 }; i += 4) {
          int inIdx = inOffset + i;
          vec4 values = vec4(
            getValue(batch, inIdx),
            getValue(batch, inIdx + 1),
            getValue(batch, inIdx + 2),
            getValue(batch, inIdx + 3)
          );

          vec4 segFilter = vec4(
            int(getSegmentIdAtIndex(inIdx)) == currentSeg ? 1 : 0,
            int(getSegmentIdAtIndex(inIdx + 1)) == currentSeg ? 1 : 0,
            int(getSegmentIdAtIndex(inIdx + 2)) == currentSeg ? 1 : 0,
            int(getSegmentIdAtIndex(inIdx + 3)) == currentSeg ? 1 : 0
          );

          ${ updateSnippet }
        }

        int inIdx = inOffset + ${ windowSizeNearestVec4 };
        if (${ windowSizeVec4Remainder === 1 }) {
          vec4 values = vec4(
            getValue(batch, inIdx),
            initializationValue,
            initializationValue,
            initializationValue
          );

          int inIdxSeg = int(getSegmentIdAtIndex(inIdx));

          vec4 segFilter = vec4(
            int(getSegmentIdAtIndex(inIdx)) == currentSeg ? 1 : 0,
            0,
            0,
            0
          );

          ${ updateSnippet }
        } else if (${ windowSizeVec4Remainder === 2 }) {
          vec4 values = vec4(
            getValue(batch, inIdx),
            getValue(batch, inIdx + 1),
            initializationValue,
            initializationValue
          );

          vec4 segFilter = vec4(
            int(getSegmentIdAtIndex(inIdx)) == currentSeg ? 1 : 0,
            int(getSegmentIdAtIndex(inIdx + 1)) == currentSeg ? 1 : 0,
              0,
              0
          );

          ${ updateSnippet }
        } else if (${ windowSizeVec4Remainder === 3 }) {
          vec4 values = vec4(
            getValue(batch, inIdx),
            getValue(batch, inIdx + 1),
            getValue(batch, inIdx + 2),
            initializationValue
          );

          vec4 segFilter = vec4(
            int(getSegmentIdAtIndex(inIdx)) == currentSeg ? 1 : 0,
            int(getSegmentIdAtIndex(inIdx + 1)) == currentSeg ? 1 : 0,
            int(getSegmentIdAtIndex(inIdx + 2)) == currentSeg ? 1 : 0,
            0
          );

          ${ updateSnippet }
        }
        setOutput(${ returnValue });
      }
    ` ; } }