PTO-BLOCK-MODEL-DISPATCH-SHARED-CUBE-MATRIX

1// PTO-UNIT: {"id":"PTO-BLOCK-MODEL-DISPATCH-SHARED-CUBE-MATRIX","surface":"block","classification":["model","dispatch","shared-cube-matrix"],"depends_on":["PTO-BLOCK-MODEL-OPERANDS-SHARED-GENERATION","PTO-TILE-MODEL-LEGALITY-MATRIX-FUNCTIONS","PTO-TILE-MODEL-STATE-SHARED-REGISTERS"]}2
3readonly func BundleMatrixSharedSourceSchemaLegal(4    ordinal: integer {0..3},5    valid_rows: integer {1..65535},6    valid_columns: integer {1..65535},7    columns: integer {1..65535},8    data_type: TileDataType) => boolean9begin10    let shared_tile_id = BundleSharedBindingId(ordinal);11    if BundleSharedBindingIsDestination(ordinal) ||12       !SharedTileCooperativeMatrixReady(shared_tile_id) then13        return FALSE;14    end;15    if _BundleSharedBindings[[ordinal]].source0_subview.valid then16        if !BundleSharedSubviewLegal(ordinal) then return FALSE; end;17        let view = MaterializeBundleSharedSubview(ordinal);18        return view.valid_rows == valid_rows &&19               view.valid_columns == valid_columns &&20               view.columns == columns && view.data_type == data_type &&21               view.layout == TileLayout_RowMajor;22    end;23    return SharedTileReadSchemaLegal(shared_tile_id,24        valid_rows, valid_columns, columns,25        data_type, TileLayout_RowMajor);26end;27
28readonly func BundleMatrixSharedSourcesReady(29    shared_count: integer {0..4}) => boolean30begin31    if shared_count == 0 then return TRUE; end;32    for ordinal = 0 to shared_count - 1 looplimit 4 do33        let binding = ordinal as integer {0..3};34        if _BundleSharedBindings[[binding]].valid &&35           !BundleSharedBindingIsDestination(binding) &&36           !SharedTilePublished(BundleSharedBindingId(binding)) then37            return FALSE;38        end;39    end;40    return TRUE;41end;42
43readonly func BundleMatrixSharedPrimarySchemaLegal(44    ordinal: integer {0..3},45    logical_rows: integer {1..65535},46    logical_columns: integer {1..65535},47    data_type: TileDataType,48    transpose: boolean) => boolean49begin50    let stored_rows = if transpose then logical_columns else logical_rows;51    let stored_columns = if transpose then logical_rows else logical_columns;52    return BundleMatrixSharedSourceSchemaLegal(53        ordinal, stored_rows, stored_columns,54        stored_columns, data_type);55end;56
57pure func BundleMatrixCooperativeMPerPE(58    group_m: integer {1..65535}) => integer {0,16,32}59begin60    if group_m <= 64 then return 16;61    elsif group_m <= 128 then return 32;62    else return 0;63    end;64end;65
66pure func BundleMatrixCooperativeValidM(67    group_m: integer {1..65535},68    pe_identity: MemoryAgentId) => integer {0..32}69begin70    let m_per_pe = BundleMatrixCooperativeMPerPE(group_m);71    if m_per_pe == 0 then return 0; end;72    let first_row = pe_identity * m_per_pe;73    if first_row >= group_m then return 0; end;74    let remaining = group_m - first_row;75    if remaining < m_per_pe then return remaining as integer {1..31}; end;76    return m_per_pe as integer {16,32};77end;78
79pure func BundleMatrixCooperativeCurrentPEMask(80    group_m: integer {1..65535},81    pe_identity: MemoryAgentId) => bits(4)82begin83    var mask = Zeros{4};84    if BundleMatrixCooperativeValidM(group_m, pe_identity) != 0 then85        mask[PTOPEMaskBitOfPEIdentity(pe_identity)] = '1';86    end;87    return mask;88end;89
90readonly func BundleMatrixSharedLeftPrimarySchemaLegal(91    ordinal: integer {0..3},92    group_m: integer {1..65535},93    k: integer {1..65535},94    data_type: TileDataType,95    transpose: boolean) => boolean96begin97    if BundleMatrixCooperativeMPerPE(group_m) == 0 then return FALSE; end;98    return BundleMatrixSharedPrimarySchemaLegal(99        ordinal, group_m, k, data_type, transpose);100end;101
102readonly func BundleMatrixSharedSchemasLegal(103    function: integer {0..31},104    left_type: TileDataType,105    right_type: TileDataType,106    m: integer {1..65535},107    n: integer {1..65535},108    k: integer {1..65535},109    shared_count: integer {0..4}) => boolean110begin111    if shared_count == 0 then112        return !_BundleFixedPointAttributes.trans_a &&113               !_BundleFixedPointAttributes.trans_b;114    end;115    let left_scale_present = TileMatrixFunctionUsesMX(function) &&116        TileMXInputTypeNeedsScale(left_type);117    let right_scale_present = TileMatrixFunctionUsesMX(function) &&118        TileMXInputTypeNeedsScale(right_type);119    let left_scale_groups = if left_scale_present then120        TileMXScaleGroupCount(k, left_type) else 1;121    let right_scale_groups = if right_scale_present then122        TileMXScaleGroupCount(k, right_type) else 1;123    let right_group = TileMatrixRightGroupSourceCount(124        function, right_type);125    let left_shared = shared_count != right_group;126    if _BundleFixedPointAttributes.trans_a && !left_shared then127        return FALSE;128    end;129    var ordinal: integer {0..4} = 0;130
131    if left_shared then132        if !BundleMatrixSharedLeftPrimarySchemaLegal(133               ordinal as integer {0..3},134               m, k, left_type,135               _BundleFixedPointAttributes.trans_a) then136            return FALSE;137        end;138        ordinal = (ordinal + 1) as integer {0..4};139        if left_scale_present then140            if !BundleMatrixSharedPrimarySchemaLegal(141                   ordinal as integer {0..3},142                   m, left_scale_groups,143                   TileMXScaleCarrierType(left_type),144                   _BundleFixedPointAttributes.trans_a) then145                return FALSE;146            end;147            ordinal = (ordinal + 1) as integer {0..4};148        end;149    end;150
151    if !BundleMatrixSharedPrimarySchemaLegal(152           ordinal as integer {0..3},153           k, n, right_type,154           _BundleFixedPointAttributes.trans_b) then155        return FALSE;156    end;157    ordinal = (ordinal + 1) as integer {0..4};158    if right_scale_present then159        if !BundleMatrixSharedPrimarySchemaLegal(160               ordinal as integer {0..3},161               right_scale_groups, n,162               TileMXScaleCarrierType(right_type),163               _BundleFixedPointAttributes.trans_b) then164            return FALSE;165        end;166        ordinal = (ordinal + 1) as integer {0..4};167    end;168    return ordinal == shared_count;169end;170
171readonly func MaterializeBundleSharedMatrixLeftPrimary(172    ordinal: integer {0..3},173    group_m: integer {1..65535},174    k: integer {1..65535},175    data_type: TileDataType,176    transpose: boolean,177    pe_identity: MemoryAgentId) => TileInfo178begin179    assert BundleMatrixSharedLeftPrimarySchemaLegal(180        ordinal, group_m, k, data_type, transpose);181    let m_per_pe = BundleMatrixCooperativeMPerPE(group_m);182    let valid_m = BundleMatrixCooperativeValidM(group_m, pe_identity);183    assert m_per_pe != 0 && valid_m != 0;184    let pe_m = valid_m as integer {1..32};185    let shared_tile_id = BundleSharedBindingId(ordinal);186    let source = if187        _BundleSharedBindings[[ordinal]].source0_subview.valid then188        MaterializeBundleSharedSubview(ordinal)189    else SharedTileRecord(shared_tile_id).tile;190    var tile = source;191    tile.contents_defined = FALSE;192    tile.defined_elements = Zeros{PTO_MODEL_TILE_ELEMENTS};193    tile.packed_defined_elements = ZeroPackedTileDefinedElements();194    tile.defined_valid_elements = 0;195    tile.rows = DerivedTileRows(tile.capacity_bytes, k, data_type);196    tile.columns = k;197    tile.valid_rows = pe_m;198    tile.valid_columns = k;199    tile.data_type = data_type;200    tile.layout = TileLayout_RowMajor;201    tile.location = TileLocation_Any;202    tile.cube_k_repeat = 0;203    tile.cube_n_repeat = 0;204    tile.cube_cell_count = 0;205    tile.cube_storage_bytes = 0;206    for row = 0 to pe_m - 1 looplimit 65536 do207        for column = 0 to k - 1 looplimit 65536 do208            let group_row = (pe_identity * m_per_pe + row)209                as integer {0..65535};210            let source_row = if transpose then column else group_row;211            let source_column = if transpose then group_row else column;212            let source_element = TileLogicalLinearIndex(213                source,214                source_row as integer {0..65535},215                source_column as integer {0..65535});216            let destination_element = TileLogicalLinearIndex(217                tile,218                row as integer {0..65535},219                column as integer {0..65535});220            tile = TileInfoWithLogicalElement(tile, destination_element,221                TileReadLogicalElement(source, source_element));222        end;223    end;224    tile.contents_defined = TRUE;225    tile.defined_valid_elements = (pe_m * k) as integer {0..524288};226    return tile;227end;228
229readonly func MaterializeBundleSharedMatrixLeftScale(230    ordinal: integer {0..3},231    group_m: integer {1..65535},232    k: integer {1..65535},233    primary_type: TileDataType,234    transpose: boolean,235    pe_identity: MemoryAgentId) => TileInfo236begin237    let m_per_pe = BundleMatrixCooperativeMPerPE(group_m);238    let valid_m = BundleMatrixCooperativeValidM(group_m, pe_identity);239    assert m_per_pe != 0 && valid_m != 0;240    let pe_m = valid_m as integer {1..32};241    let scale_groups = TileMXScaleGroupCount(k, primary_type);242    let scale_type = TileMXScaleCarrierType(primary_type);243    let shared_tile_id = BundleSharedBindingId(ordinal);244    let source = if245        _BundleSharedBindings[[ordinal]].source0_subview.valid then246        MaterializeBundleSharedSubview(ordinal)247    else SharedTileRecord(shared_tile_id).tile;248    assert BundleMatrixSharedPrimarySchemaLegal(249        ordinal, group_m, scale_groups, scale_type, transpose);250    var tile = source;251    tile.contents_defined = FALSE;252    tile.defined_elements = Zeros{PTO_MODEL_TILE_ELEMENTS};253    tile.packed_defined_elements = ZeroPackedTileDefinedElements();254    tile.defined_valid_elements = 0;255    tile.rows = DerivedTileRows(256        tile.capacity_bytes, scale_groups, scale_type);257    tile.columns = scale_groups;258    tile.valid_rows = pe_m;259    tile.valid_columns = scale_groups;260    tile.data_type = scale_type;261    tile.layout = TileLayout_RowMajor;262    tile.location = TileLocation_Any;263    tile.cube_k_repeat = 0;264    tile.cube_n_repeat = 0;265    tile.cube_cell_count = 0;266    tile.cube_storage_bytes = 0;267    for row = 0 to pe_m - 1 looplimit 65536 do268        for column = 0 to scale_groups - 1 looplimit 2048 do269            let group_row = (pe_identity * m_per_pe + row)270                as integer {0..65535};271            let source_row = if transpose then column else group_row;272            let source_column = if transpose then group_row else column;273            let source_element = TileLogicalLinearIndex(274                source,275                source_row as integer {0..65535},276                source_column as integer {0..65535});277            let destination_element = TileLogicalLinearIndex(278                tile, row as integer {0..65535},279                column as integer {0..65535});280            tile = TileInfoWithLogicalElement(tile, destination_element,281                TileReadLogicalElement(source, source_element));282        end;283    end;284    tile.contents_defined = TRUE;285    tile.defined_valid_elements =286        (pe_m * scale_groups) as integer {0..524288};287    return tile;288end;289
290readonly func MaterializeBundleSharedMatrixSource(291    ordinal: integer {0..3},292    valid_rows: integer {1..65535},293    valid_columns: integer {1..65535},294    columns: integer {1..65535},295    data_type: TileDataType) => TileInfo296begin297    let shared_tile_id = BundleSharedBindingId(ordinal);298    var tile = if299        _BundleSharedBindings[[ordinal]].source0_subview.valid then300        MaterializeBundleSharedSubview(ordinal)301    else MaterializeSharedTileForReadSchema(302        shared_tile_id, valid_rows, valid_columns, columns,303        data_type, TileLayout_RowMajor);304    for element = 0 to tile.rows * tile.columns - 1305        looplimit 524288 do306        let index = element as PackedTileElementIndex;307        tile = TileInfoWithLogicalElement(tile, index,308            TileReadLogicalElement(tile, index));309    end;310    tile.contents_defined = TRUE;311    tile.defined_valid_elements =312        (valid_rows * valid_columns) as integer {0..524288};313    return tile;314end;315
316readonly func MaterializeBundleSharedMatrixPrimary(317    ordinal: integer {0..3},318    logical_rows: integer {1..65535},319    logical_columns: integer {1..65535},320    data_type: TileDataType,321    transpose: boolean,322    pe_identity: MemoryAgentId) => TileInfo323begin324    assert BundleMatrixSharedPrimarySchemaLegal(325        ordinal, logical_rows, logical_columns, data_type, transpose);326    let shared_tile_id = BundleSharedBindingId(ordinal);327    let source = if328        _BundleSharedBindings[[ordinal]].source0_subview.valid then329        MaterializeBundleSharedSubviewForPE(ordinal, pe_identity)330    else SharedTileRecord(shared_tile_id).tile;331    var tile = source;332    tile.contents_defined = FALSE;333    tile.defined_elements = Zeros{PTO_MODEL_TILE_ELEMENTS};334    tile.packed_defined_elements = ZeroPackedTileDefinedElements();335    tile.defined_valid_elements = 0;336    tile.rows = DerivedTileRows(337        tile.capacity_bytes, logical_columns, data_type);338    tile.columns = logical_columns;339    tile.valid_rows = logical_rows;340    tile.valid_columns = logical_columns;341    tile.data_type = data_type;342    tile.layout = TileLayout_RowMajor;343    tile.location = TileLocation_Any;344    tile.cube_k_repeat = 0;345    tile.cube_n_repeat = 0;346    tile.cube_cell_count = 0;347    tile.cube_storage_bytes = 0;348    for row = 0 to logical_rows - 1 looplimit 65536 do349        for column = 0 to logical_columns - 1 looplimit 65536 do350            let source_row = if transpose then column else row;351            let source_column = if transpose then row else column;352            let source_element = TileLogicalLinearIndex(353                source,354                source_row as integer {0..65535},355                source_column as integer {0..65535});356            let destination_element = TileLogicalLinearIndex(357                tile,358                row as integer {0..65535},359                column as integer {0..65535});360            tile = TileInfoWithLogicalElement(tile, destination_element,361                TileReadLogicalElement(source, source_element));362        end;363    end;364    tile.contents_defined = TRUE;365    tile.defined_valid_elements =366        (logical_rows * logical_columns) as integer {0..524288};367    return tile;368end;369
370readonly func BundleMatrixCooperativeMLayout(371    function: integer {0..31},372    right_type: TileDataType,373    m: integer {1..65535},374    shared_count: integer {0..4}) => (boolean, TileLayout)375begin376    let right_group = TileMatrixRightGroupSourceCount(377        function, right_type);378    let local_left_present = shared_count == 0 ||379        shared_count == right_group;380    var layout = TileLayout_RowMajor;381    if local_left_present then382        let left_ordinal = if TileMatrixFunctionUsesAccumulator(function)383            then 1 else 0;384        layout = _Tiles[[BundleMatrixSourceAt(385            left_ordinal as integer {0..7})]].layout;386    elsif TileMatrixFunctionUsesAccumulator(function) then387        layout = _Tiles[[BundleMatrixSourceAt(0)]].layout;388    elsif m <= 16 then389        layout = TileLayout_CUBE_M16;390    elsif m <= 32 then391        layout = TileLayout_CUBE_M32;392    end;393    return (TileMatrixMLayoutLegal(layout, m), layout);394end;395