PTO-TILE-MODEL-LEGALITY-MATRIX-OPERANDS

1// PTO-UNIT: {"id":"PTO-TILE-MODEL-LEGALITY-MATRIX-OPERANDS","surface":"tile","classification":["model","legality","matrix-operands"],"depends_on":["PTO-TILE-MODEL-LEGALITY-MATRIX-CUBE-PRIMARY","PTO-TILE-MODEL-LEGALITY-MATRIX-POSTPROCESS"]}2// PTO-REQ-CUBE-OPERANDS-001: every Local matrix source descriptor is checked3// in stream order before any Local or Shared payload is snapshotted.4
5readonly func TileMatrixLocalOperandSchemaLegal(6    source: TileIndex,7    valid_rows: integer {1..65535},8    valid_columns: integer {1..65535},9    data_type: TileDataType) => boolean10begin11    let tile = _Tiles[[source]];12    return TileSourceContentsDefined(source) &&13           IsNonzeroPowerOfTwo(tile.rows) &&14           IsNonzeroPowerOfTwo(tile.columns) &&15           tile.valid_rows == valid_rows &&16           tile.valid_columns == valid_columns &&17           tile.data_type == data_type &&18           tile.layout == TileLayout_RowMajor &&19           tile.location == TileLocation_Matrix;20end;21
22readonly func TileMatrixLocalAScaleSchemaLegal(23    source: TileIndex,24    valid_rows: integer {1..65535},25    valid_columns: integer {1..65535},26    primary_type: TileDataType) => boolean27begin28    let tile = _Tiles[[source]];29    return tile.contents_defined && TileCubeDescriptorLegal(tile) &&30           tile.valid_rows == valid_rows &&31           tile.valid_columns == valid_columns &&32           tile.data_type == TileMXScaleCarrierType(primary_type) &&33           tile.layout == TileLayout_CUBE_M32 &&34           tile.location == TileLocation_Matrix;35end;36
37readonly func TileMatrixLocalBScaleSchemaLegal(38    source: TileIndex,39    groups: integer {1..65535},40    n: integer {1..65535},41    primary_type: TileDataType) => boolean42begin43    let tile = _Tiles[[source]];44    return tile.contents_defined && TileCubeDescriptorLegal(tile) &&45           tile.valid_rows == n &&46           tile.valid_columns == groups &&47           tile.data_type == TileMXScaleCarrierType(primary_type) &&48           tile.layout == TileLayout_CUBE_M32 &&49           tile.location == TileLocation_Matrix;50end;51
52readonly func TileMatrixLocalBiasSchemaLegal(53    source: TileIndex,54    n: integer {1..65535},55    accumulator_type: TileDataType) => boolean56begin57    let tile = _Tiles[[source]];58    return TileSourceContentsDefined(source) &&59           tile.valid_rows == 1 &&60           tile.valid_columns == n &&61           tile.data_type == accumulator_type &&62           tile.layout == TileLayout_RowMajor;63end;64
65readonly func TileMatrixLocalCScaleSchemaLegal(66    source: TileIndex,67    m: integer {1..65535}) => boolean68begin69    let tile = _Tiles[[source]];70    return tile.contents_defined && TileCubeDescriptorLegal(tile) &&71           tile.valid_rows == m && tile.valid_columns == 1 &&72           tile.data_type == TileDataType_U8 &&73           tile.layout == TileLayout_CUBE_M32 &&74           tile.location == TileLocation_Matrix;75end;76
77readonly func TileMatrixInfoAccumulatorSchemaLegal(78    accumulator: TileIndex,79    m: integer {1..65535},80    n: integer {1..65535},81    result_type: TileDataType,82    destination_capacity: integer {0..262144}) => boolean83begin84    let tile = _Tiles[[accumulator]];85    let output_converted =86        UInt(_BundleFixedPointAttributes.pre_quant_mode) != 0;87    return TileSourceContentsDefined(accumulator) &&88           TileInfoDescriptorLegal(tile) &&89           tile.valid_rows == m &&90           tile.valid_columns == n &&91           tile.data_type == result_type &&92           tile.layout == TileLayout_RowMajor &&93           tile.location == TileLocation_Matrix &&94           (output_converted ||95            tile.capacity_bytes == destination_capacity);96end;97
98readonly func BundleMatrixLocalMathematicalSourcesLegal(99    function: integer {0..31},100    left_type: TileDataType,101    right_type: TileDataType,102    m: integer {1..65535},103    n: integer {1..65535},104    k: integer {1..65535},105    shared_count: integer {0..4},106    accumulator_type: TileDataType,107    destination_capacity: integer {0..262144}) => boolean108begin109    let left_scale_present = TileMatrixFunctionUsesMX(function) &&110        TileMXInputTypeNeedsScale(left_type);111    let right_scale_present = TileMatrixFunctionUsesMX(function) &&112        TileMXInputTypeNeedsScale(right_type);113    let left_scale_groups = if left_scale_present then114        TileMXScaleGroupCount(k, left_type) else 1;115    let right_scale_groups = if right_scale_present then116        TileMXScaleGroupCount(k, right_type) else 1;117    var ordinal: integer {0..6} = 0;118    let right_group = TileMatrixRightGroupSourceCount(119        function, right_type);120    let local_left_present = shared_count == 0 ||121        shared_count == right_group;122    let left_ordinal = if TileMatrixFunctionUsesAccumulator(function)123        then 1 else 0;124    let local_m_layout = if local_left_present then125        _Tiles[[BundleMatrixSourceAt(126            left_ordinal as integer {0..8})]].layout127        else if TileMatrixFunctionUsesAccumulator(function) then128            _Tiles[[BundleMatrixSourceAt(0)]].layout129        else TileLayout_RowMajor;130
131    if TileMatrixFunctionUsesAccumulator(function) then132        let accumulator = BundleMatrixSourceAt(133            ordinal as integer {0..8});134        let accumulator_legal = TileMatrixLocalCubeAccumulatorSchemaLegal(135            accumulator, m, n, accumulator_type,136            local_m_layout, destination_capacity);137        if !accumulator_legal then138            return FALSE;139        end;140        ordinal = (ordinal + 1) as integer {0..6};141    end;142
143    if shared_count == 0 || shared_count == right_group then144        let left = BundleMatrixSourceAt(145            ordinal as integer {0..8});146        let left_legal = TileMatrixLocalMOperandSchemaLegal(147            left, m, k, left_type);148        if !left_legal then149            return FALSE;150        end;151        ordinal = (ordinal + 1) as integer {0..6};152        if left_scale_present then153            let left_scale = BundleMatrixSourceAt(154                ordinal as integer {0..8});155            if !TileMatrixLocalAScaleSchemaLegal(156                   left_scale, m, left_scale_groups, left_type) then157                return FALSE;158            end;159            ordinal = (ordinal + 1) as integer {0..6};160        end;161    end;162
163    if shared_count == 0 then164        let right = BundleMatrixSourceAt(165            ordinal as integer {0..8});166        let right_legal = TileMatrixLocalNOperandSchemaLegal(167            right, k, n, right_type);168        if !right_legal then169            return FALSE;170        end;171        ordinal = (ordinal + 1) as integer {0..6};172        if right_scale_present then173            let right_scale = BundleMatrixSourceAt(174                ordinal as integer {0..8});175            if !TileMatrixLocalBScaleSchemaLegal(176                   right_scale, right_scale_groups, n, right_type) then177                return FALSE;178            end;179            ordinal = (ordinal + 1) as integer {0..6};180        end;181    end;182
183    if TileMatrixFunctionUsesBias(function) then184        let bias = BundleMatrixSourceAt(185            ordinal as integer {0..8});186        if !TileMatrixLocalBiasSchemaLegal(187               bias, n, accumulator_type) then188            return FALSE;189        end;190        ordinal = (ordinal + 1) as integer {0..6};191    end;192
193    if _BundleFixedPointAttributes.c_scale_en then194        let c_scale = BundleMatrixSourceAt(195            ordinal as integer {0..8});196        if !TileMatrixFunctionAllowsCScale(function) ||197           accumulator_type != TileDataType_FP32 ||198           !TileMatrixLocalCScaleSchemaLegal(c_scale, m) then199            return FALSE;200        end;201    end;202    return TRUE;203end;204