PTO-TILE-MODEL-EXECUTION-EXPANSION

1// PTO-UNIT: {"id":"PTO-TILE-MODEL-EXECUTION-EXPANSION","surface":"tile","classification":["model","execution","expansion"],"depends_on":["PTO-TILE-MODEL-EXECUTION-REDUCTION","PTO-TILE-MODEL-EXECUTION-UNARY"]}2// PTO-REQ-TEPL-EXPAND-001: exact typed row and column broadcast operations.3
4pure func TileExpandBinaryOperation(5    operation: TileExpandOperation) => TileBinaryOperation6begin7    case operation of8        when TileExpand_ADD =>9            return TileBinary_ADD;10        when TileExpand_SUB =>11            return TileBinary_SUB;12        when TileExpand_MUL =>13            return TileBinary_MUL;14        when TileExpand_DIV =>15            return TileBinary_DIV;16        when TileExpand_MAX =>17            return TileBinary_MAX;18        when TileExpand_MIN =>19            return TileBinary_MIN;20        otherwise =>21            unreachable;22    end;23end;24
25// Mixed EXPDIF owns these conversions locally.  They are exact value26// widenings for the selected source formats; they are not TCVT and never27// produce a conversion-inexact status.28pure func ExactWidenBF16ToFP32(value: Word) => Word29begin30    return LSL(ZeroExtend{PTO_XLEN}(value[15:0]), 16);31end;32
33pure func ExactWidenFP16ToFP32(value: Word) => Word34begin35    let raw = value[15:0];36    let sign = raw[15];37    let exponent = raw[14:10];38    let fraction = raw[9:0];39    var result: bits(32) = Zeros{32};40    result[31] = sign;41    if exponent == '11111' then42        result[30:23] = Ones{8};43        if fraction != Zeros{10} then44            // The selected IEEE profile canonicalizes produced FP32 NaNs.45            result[31] = '0';46            result[22:0] = Zeros{23} + 0x400000;47        end;48    elsif exponent != Zeros{5} then49        result[30:23] = Zeros{8} + (UInt(exponent) + 112);50        result[22:13] = fraction;51    elsif fraction != Zeros{10} then52        var normalized = fraction;53        var shift_count: integer {0..9} = 0;54        for shift = 0 to 9 looplimit 10 do55            if normalized[9] == '0' then56                normalized = LSL(normalized, 1);57                shift_count = (shift_count + 1) as integer {0..9};58            end;59        end;60        result[30:23] = Zeros{8} + (112 - shift_count);61        result[22:13] = ZeroExtend{10}(normalized[8:0]);62    end;63    return ZeroExtend{PTO_XLEN}(result);64end;65
66func TileExpandValueWithTypesAndFlags(67    operation: TileExpandOperation,68    source_type: TileDataType,69    destination_type: TileDataType,70    left: Word,71    broadcast: Word) => (Word, bits(5))72begin73    if operation == TileExpand_COPY then74        return (broadcast, Zeros{5});75    end;76
77    if operation == TileExpand_EXPDIF then78        if source_type != destination_type then79            assert destination_type == TileDataType_FP32;80            let widened_left = if source_type == TileDataType_FP16 then81                ExactWidenFP16ToFP32(left)82            else83                ExactWidenBF16ToFP32(left);84            let widened_broadcast = if source_type == TileDataType_FP16 then85                ExactWidenFP16ToFP32(broadcast)86            else87                ExactWidenBF16ToFP32(broadcast);88            // The named IEEE profile owns FP32 SUB/EXP for this mixed89            // path.  The operation-local widening above remains portable.90            let (profile_result, profile_flags) =91                TileProfileMixedExpdifFP32(92                    source_type,93                    widened_left,94                    widened_broadcast);95            return (96                profile_result,97                profile_flags);98        end;99        let (difference, subtract_flags) =100            TileProfileBinaryWithFlags(101                TileBinary_SUB,102                destination_type,103                left,104                broadcast);105        let (handled, special_result, special_flags) =106            TileSFUUnarySpecialValue(107                TileUnary_EXP,108                destination_type,109                difference);110        if handled then111            return (112                special_result,113                subtract_flags OR special_flags);114        end;115        let (profile_result, profile_flags) = TileProfileUnary(116            TileUnary_EXP,117            destination_type,118            difference);119        return (120            profile_result,121            subtract_flags OR profile_flags);122    end;123
124    return TileProfileBinaryWithFlags(125        TileExpandBinaryOperation(operation),126        destination_type,127        left,128        broadcast);129end;130
131impdef func TileProfileMixedExpdifFP32(132    source_type: TileDataType,133    left: Word,134    broadcast: Word) => (Word, bits(5))135begin136    assert source_type == TileDataType_FP16 ||137           source_type == TileDataType_BF16;138    let (difference, subtract_flags) = TileProfileBinaryWithFlags(139        TileBinary_SUB,140        TileDataType_FP32,141        left,142        broadcast);143    let (handled, special_result, special_flags) = TileSFUUnarySpecialValue(144        TileUnary_EXP,145        TileDataType_FP32,146        difference);147    if handled then148        return (149            special_result,150            subtract_flags OR special_flags);151    end;152    let (profile_result, profile_flags) = TileProfileUnary(153        TileUnary_EXP,154        TileDataType_FP32,155        difference);156    return (157        profile_result,158        subtract_flags OR profile_flags);159end;160
161impdef func TileProfileExpand(op: TileExpandOperation,162                              data_type: TileDataType,163                              left: Word, broadcast: Word) => Word164begin165    return TileExpandValue(166        op,167        data_type,168        left,169        broadcast);170end;171
172func TileExpandValueWithFlags(173    operation: TileExpandOperation,174    data_type: TileDataType,175    left: Word,176    broadcast: Word) => (Word, bits(5))177begin178    return TileExpandValueWithTypesAndFlags(179        operation,180        data_type,181        data_type,182        left,183        broadcast);184end;185
186func TileExpandValue(187    operation: TileExpandOperation,188    data_type: TileDataType,189    left: Word,190    broadcast: Word) => Word191begin192    let (result, -) = TileExpandValueWithFlags(193        operation,194        data_type,195        left,196        broadcast);197    return result;198end;199
200func ExecuteTileExpand(op: TileExpandOperation, axis: TileAxis,201                       destination: TileIndex, source: TileIndex,202                       broadcast_source: TileIndex)203begin204    assert TileOperandsLegal_ExecuteTileExpand(205        op,206        axis,207        destination,208        source,209        broadcast_source);210
211    let source_tile = _Tiles[[source]];212    let broadcast_tile = _Tiles[[broadcast_source]];213    var result_tile = _Tiles[[destination]];214    var accumulated_flags = Zeros{5};215
216    for row = 0 to result_tile.valid_rows - 1 looplimit 65536 do217        for column = 0 to result_tile.valid_columns - 1 looplimit 65536 do218            let broadcast_row = if axis == TileAxis_Row then row else 0;219            let broadcast_column = if axis == TileAxis_Row then 0 else column;220            let broadcast_element = TileLogicalLinearIndex(broadcast_tile,221                broadcast_row as integer {0..65535},222                broadcast_column as integer {0..65535});223            var left = TileReadLogicalElement(broadcast_tile,224                broadcast_element);225            if op != TileExpand_COPY then226                let source_element = TileLogicalLinearIndex(227                    source_tile,228                    row as integer {0..65535},229                    column as integer {0..65535});230                left = TileReadLogicalElement(source_tile, source_element);231            end;232            let (value, element_flags) = TileExpandValueWithTypesAndFlags(233                op,234                source_tile.data_type,235                result_tile.data_type,236                left,237                TileReadLogicalElement(broadcast_tile, broadcast_element));238            let destination_element = TileLogicalLinearIndex(239                result_tile,240                row as integer {0..65535},241                column as integer {0..65535});242            result_tile = TileInfoWithLogicalElement(result_tile,243                destination_element, value);244            accumulated_flags = accumulated_flags OR element_flags;245        end;246    end;247
248    result_tile = TileWithValidRegionDefined(result_tile);249    result_tile = TileWithPadding(250        result_tile,251        CurrentBundlePadValue());252    RecordNumericStatusFlags(accumulated_flags);253    _Tiles[[destination]] = result_tile;254end;255