PTO-TILE-MODEL-LEGALITY-MATRIX-FUNCTIONS

1// PTO-UNIT: {"id":"PTO-TILE-MODEL-LEGALITY-MATRIX-FUNCTIONS","surface":"tile","classification":["model","legality","matrix-functions"],"depends_on":["PTO-ARCH-DATA-TYPES-TILE-DATA-TYPES"]}2// The CUBE Matrix selector and MX side-type rules are kept in one small unit3// so schema and execution code cannot grow separate function tables.4
5// NDF-BEGIN: PTO-CUBE-MATRIX-SCALE-0016// ndf: kind=contract level=L1 layer=tile status=accepted7// Each Matrix-MX primary side MUST independently select group-32 E8M0 scale8// for MX FP8/FP4 carriers or group-64 raw U32 scale for HiF4X2. HiF4X2 MUST9// be accepted only by Matrix-MX input roles; ordinary Matrix MUST not gain it.10// Each Local scale MUST use CUBE_M32 and each Shared scale MUST remain an11// independently bound ordinary Tile with the corresponding primary location.12// NDF-END: PTO-CUBE-MATRIX-SCALE-00113
14pure func TileMXInputTypeSupported(data_type: TileDataType) => boolean15begin16    return data_type == TileDataType_FP16 ||17           data_type == TileDataType_BF16 ||18           data_type == TileDataType_E4M3 ||19           data_type == TileDataType_E5M2 ||20           data_type == TileDataType_E2M1X2 ||21           data_type == TileDataType_E1M2X2 ||22           data_type == TileDataType_HiF4X2;23end;24
25pure func TileMXScaleGroupSize(data_type: TileDataType)26    => integer {32,64}27begin28    assert TileMXInputTypeNeedsScale(data_type);29    if data_type == TileDataType_HiF4X2 then return 64; end;30    return 32;31end;32
33pure func TileMXScaleCarrierType(data_type: TileDataType) => TileDataType34begin35    assert TileMXInputTypeNeedsScale(data_type);36    if data_type == TileDataType_HiF4X2 then return TileDataType_U32; end;37    return TileDataType_E8M0;38end;39
40pure func TileMXScaleGroupCount(41    k: integer {1..65535}, data_type: TileDataType)42    => integer {1..2048}43begin44    let group_size = TileMXScaleGroupSize(data_type);45    return ((k + (group_size - 1)) DIVRM group_size)46        as integer {1..2048};47end;48
49pure func TileMXInputTypeNeedsScale(data_type: TileDataType) => boolean50begin51    assert TileMXInputTypeSupported(data_type);52    return data_type != TileDataType_FP16 &&53           data_type != TileDataType_BF16;54end;55
56pure func TileMXOperandPairLegal(left_type: TileDataType,57                                right_type: TileDataType) => boolean58begin59    return TileMXInputTypeSupported(left_type) &&60           TileMXInputTypeSupported(right_type);61end;62
63pure func TileMatrixFunctionAssigned(function: integer {0..31}) => boolean64begin65    return function == 0 || function == 1 || function == 2 ||66           function == 4 || function == 5 || function == 6 ||67           function == 16 || function == 17 || function == 18 ||68           function == 20 || function == 21 || function == 22;69end;70
71pure func TileMatrixFunctionUsesBias(function: integer {0..31}) => boolean72begin73    return function == 1 || function == 5 ||74           function == 17 || function == 21;75end;76
77pure func TileMatrixFunctionUsesAccumulator(78    function: integer {0..31}) => boolean79begin80    return function == 2 || function == 6 ||81           function == 18 || function == 22;82end;83
84pure func TileMatrixFunctionUsesMX(function: integer {0..31}) => boolean85begin86    return function == 4 || function == 5 || function == 6 ||87           function == 20 || function == 21 || function == 22;88end;89
90pure func TileMatrixFunctionIsGEMV(function: integer {0..31}) => boolean91begin92    return function == 16 || function == 17 || function == 18 ||93           function == 20 || function == 21 || function == 22;94end;95
96pure func TileMatrixFunctionAllowsCScale(97    function: integer {0..31}) => boolean98begin99    return function == 2 || function == 6;100end;101
102pure func TileMatrixLeftGroupSourceCount(103    function: integer {0..31}, left_type: TileDataType) => integer {1..2}104begin105    if TileMatrixFunctionUsesMX(function) &&106       TileMXInputTypeNeedsScale(left_type) then107        return 2;108    end;109    return 1;110end;111
112pure func TileMatrixRightGroupSourceCount(113    function: integer {0..31}, right_type: TileDataType) => integer {1..2}114begin115    if TileMatrixFunctionUsesMX(function) &&116       TileMXInputTypeNeedsScale(right_type) then117        return 2;118    end;119    return 1;120end;121
122pure func TileMatrixMathematicalSourceCount(123    function: integer {0..31}, left_type: TileDataType,124    right_type: TileDataType) => integer {2..5}125begin126    assert TileMatrixFunctionAssigned(function);127    let matrix_sources = TileMatrixLeftGroupSourceCount(128        function, left_type) + TileMatrixRightGroupSourceCount(129        function, right_type);130    let supplementary_source =131        TileMatrixFunctionUsesBias(function) ||132        TileMatrixFunctionUsesAccumulator(function);133    return (matrix_sources + (if supplementary_source then 1 else 0))134        as integer {2..5};135end;136
137// Shared Matrix inputs are carried in complete operand groups.  A right-only138// stream contains the right matrix and its optional scale.  A both-sides139// stream contains the complete left group followed by the complete right140// group.  TGEMV remains Local-only.141pure func TileMatrixSharedSourceCountLegal(142    function: integer {0..31}, left_type: TileDataType,143    right_type: TileDataType, shared_count: integer {0..4}) => boolean144begin145    assert TileMatrixFunctionAssigned(function);146    if TileMatrixFunctionIsGEMV(function) then147        return shared_count == 0;148    end;149    if shared_count == 0 then150        return TRUE;151    end;152    let right_group = TileMatrixRightGroupSourceCount(153        function, right_type);154    let both_groups = TileMatrixLeftGroupSourceCount(155        function, left_type) + right_group;156    return shared_count == right_group ||157           shared_count == both_groups;158end;159
160// Local mathematical operands preserve their architectural order after the161// Shared groups are removed.  ACC contributes C first; bias contributes the162// final Local source.  Post-processing sources are not counted here.163pure func TileMatrixLocalMathematicalSourceCount(164    function: integer {0..31}, left_type: TileDataType,165    right_type: TileDataType, shared_count: integer {0..4})166    => integer {0..5}167begin168    assert TileMatrixSharedSourceCountLegal(169        function, left_type, right_type, shared_count);170    let supplementary = if171        TileMatrixFunctionUsesBias(function) ||172        TileMatrixFunctionUsesAccumulator(function)173    then 1 else 0;174    if shared_count == 0 then175        return TileMatrixMathematicalSourceCount(176            function, left_type, right_type) as integer {0..5};177    end;178    let right_group = TileMatrixRightGroupSourceCount(179        function, right_type);180    if shared_count == right_group then181        return (TileMatrixLeftGroupSourceCount(function, left_type) +182                supplementary) as integer {0..5};183    end;184    return supplementary as integer {0..5};185end;186