license import { GPGPUProgram } from './gpgpu_math' ; import { getCoordsDataType } from './shader_compiler' ; export class ScatterProgram implements GPGPUProgram { variableNames = [ 'updates' , 'indices' , 'defaultValue' ] ; outputShape : number [ ] ; userCode : string ; constructor ( updateSize , sliceDim , indicesRank , updatesRank , strides , shape , summingDupeIndex = true ) { this . outputShape = shape ; const stridesType = getCoordsDataType ( strides . length ) ; const dtype = getCoordsDataType ( shape . length ) ; let indicesString = '' ; if ( indicesRank === 1 ) { indicesString = 'i' ; } else if ( indicesRank === 2 ) { indicesString = 'i, j' ; } const indicesSnippet = `getIndices(${ indicesString })` ; let updatesString = '' ; if ( updatesRank === 1 ) { updatesString = 'i' ; } else if ( updatesRank === 2 ) { updatesString = 'i, coords[1]' ; } const updatesSnippet = `getUpdates(${ updatesString })` ; const strideString = sliceDim > 1 ? 'strides[j]' : 'strides' ; this . userCode = `
        ${ stridesType } strides = ${ stridesType }(${ strides });

        void main() {
          ${ dtype } coords = getOutputCoords();
          float sum = 0.0;
          bool found = false;
          for (int i = 0; i < ${ updateSize }; i++) {
            int flattenedIndex = 0;
            for (int j = 0; j < ${ sliceDim }; j++) {
              int index = round(${ indicesSnippet });
              flattenedIndex += index * ${ strideString };
            }
            if (flattenedIndex == coords[0]) {
              sum += ${ updatesSnippet };
              found = true;
            }
          }
          setOutput(mix(getDefaultValue(), sum, float(found)));
        }
      ` ; } }