PTO-TILE-MODEL-EXECUTION-REDUCTION

1// PTO-UNIT: {"id":"PTO-TILE-MODEL-EXECUTION-REDUCTION","surface":"tile","classification":["model","execution","reduction"],"depends_on":["PTO-TILE-MODEL-EXECUTION-ELEMENTWISE"]}2// PTO-REQ-TEPL-REDUCE-001: exact row, column, and index reductions.3
4pure func TileReductionOneEncoding(5    data_type: TileDataType) => Word6begin7    case data_type of8        when TileDataType_FP64 =>9            return Zeros{PTO_XLEN} + 0x3ff0000000000000;10        when TileDataType_FP32, TileDataType_TF32,11             TileDataType_HF32 =>12            return Zeros{PTO_XLEN} + 0x3f800000;13        when TileDataType_FP16 =>14            return Zeros{PTO_XLEN} + 0x3c00;15        when TileDataType_BF16 =>16            return Zeros{PTO_XLEN} + 0x3f80;17        when TileDataType_E4M3 =>18            return Zeros{PTO_XLEN} + 0x38;19        when TileDataType_E5M2 =>20            return Zeros{PTO_XLEN} + 0x3c;21        when TileDataType_S64, TileDataType_S32,22             TileDataType_S16, TileDataType_S8,23             TileDataType_U64, TileDataType_U32,24             TileDataType_U16, TileDataType_U8 =>25            return Zeros{PTO_XLEN} + 1;26        otherwise =>27            unreachable;28    end;29end;30
31pure func TileReductionInitialValue(32    operation: TileReductionOperation,33    data_type: TileDataType,34    first: Word) => Word35begin36    case operation of37        when TileReduction_SUM =>38            return Zeros{PTO_XLEN};39        when TileReduction_PRODUCT =>40            return TileReductionOneEncoding(data_type);41        when TileReduction_MIN, TileReduction_MAX,42             TileReduction_ARGMIN, TileReduction_ARGMAX =>43            return first;44    end;45end;46
47impdef func TileProfileReductionInitial(48    operation: TileReductionOperation,49    data_type: TileDataType,50    first: Word) => Word51begin52    return TileReductionInitialValue(53        operation,54        data_type,55        first);56end;57
58func TileReductionStepWithFlags(59    operation: TileReductionOperation,60    data_type: TileDataType,61    accumulator: Word,62    value: Word) => (Word, boolean, bits(5))63begin64    var binary_operation: TileBinaryOperation;65    case operation of66        when TileReduction_SUM =>67            binary_operation = TileBinary_ADD;68        when TileReduction_PRODUCT =>69            binary_operation = TileBinary_MUL;70        when TileReduction_MIN, TileReduction_ARGMIN =>71            binary_operation = TileBinary_MIN;72        when TileReduction_MAX, TileReduction_ARGMAX =>73            binary_operation = TileBinary_MAX;74    end;75
76    let (result, flags) = TileProfileBinaryWithFlags(77        binary_operation,78        data_type,79        accumulator,80        value);81    let selected =82        result == value && result != accumulator;83    return (result, selected, flags);84end;85
86impdef func TileProfileReductionStep(87    operation: TileReductionOperation,88    data_type: TileDataType,89    accumulator: Word,90    value: Word) => (Word, boolean)91begin92    let (result, selected, -) = TileReductionStepWithFlags(93        operation,94        data_type,95        accumulator,96        value);97    return (result, selected);98end;99
100func ExecuteTileReduction(101    operation: TileReductionOperation,102    axis: TileAxis,103    destination: TileIndex,104    source: TileIndex)105begin106    assert TileOperandsLegal_ExecuteTileReduction(107        operation,108        axis,109        destination,110        source);111
112    let source_tile = _Tiles[[source]];113    var result_tile = _Tiles[[destination]];114    var accumulated_flags = Zeros{5};115    let outer_count =116        if axis == TileAxis_Row then117            source_tile.valid_rows118        else119            source_tile.valid_columns;120    let inner_count =121        if axis == TileAxis_Row then122            source_tile.valid_columns123        else124            source_tile.valid_rows;125
126    for outer = 0 to outer_count - 1 looplimit 65536 do127        let first_row =128            if axis == TileAxis_Row then outer else 0;129        let first_column =130            if axis == TileAxis_Row then 0 else outer;131        let first_element = TileLogicalLinearIndex(132            source_tile,133            first_row as integer {0..65535},134            first_column as integer {0..65535});135        var accumulator = TileProfileReductionInitial(136            operation,137            source_tile.data_type,138            TileReadLogicalElement(source_tile, first_element));139        var selected_index: integer {0..65535} = 0;140        let identity_reduction =141            operation == TileReduction_SUM ||142            operation == TileReduction_PRODUCT;143        let first_inner =144            if identity_reduction then 0 else 1;145
146        if first_inner < inner_count then147            for inner = first_inner to inner_count - 1148                looplimit 65536 do149                let row =150                    if axis == TileAxis_Row then outer else inner;151                let column =152                    if axis == TileAxis_Row then inner else outer;153                let element = TileLogicalLinearIndex(154                    source_tile,155                    row as integer {0..65535},156                    column as integer {0..65535});157                let (next, selected, element_flags) =158                    TileReductionStepWithFlags(159                        operation,160                        source_tile.data_type,161                        accumulator,162                        TileReadLogicalElement(source_tile, element));163                accumulator = next;164                accumulated_flags =165                    accumulated_flags OR element_flags;166                if selected &&167                   (operation == TileReduction_ARGMIN ||168                    operation == TileReduction_ARGMAX) then169                    selected_index =170                        inner as integer {0..65535};171                end;172            end;173        end;174
175        let destination_row =176            if axis == TileAxis_Row then outer else 0;177        let destination_column =178            if axis == TileAxis_Row then 0 else outer;179        let destination_element = TileLogicalLinearIndex(180            result_tile,181            destination_row as integer {0..65535},182            destination_column as integer {0..65535});183        if operation == TileReduction_ARGMIN ||184           operation == TileReduction_ARGMAX then185            result_tile = TileInfoWithLogicalElement(result_tile,186                destination_element, NaturalToWord(187                    selected_index as integer {0..262144}));188        else189            result_tile = TileInfoWithLogicalElement(result_tile,190                destination_element, accumulator);191        end;192    end;193
194    result_tile = TileWithValidRegionDefined(result_tile);195    result_tile = TileWithPadding(196        result_tile,197        CurrentBundlePadValue());198    RecordNumericStatusFlags(accumulated_flags);199    _Tiles[[destination]] = result_tile;200end;201