PTO-BLOCK-MODEL-DISPATCH-CUBE-DESTINATION

1// PTO-UNIT: {"id":"PTO-BLOCK-MODEL-DISPATCH-CUBE-DESTINATION","surface":"block","classification":["model","dispatch","cube-destination"],"depends_on":["PTO-BLOCK-MODEL-DISPATCH-DESTINATION-SHAPE","PTO-BLOCK-MODEL-FAULTS-ROLLBACK","PTO-TILE-MODEL-LEGALITY-MATRIX-CUBE-PRIMARY"]}2
3readonly func BundleTMATMULDestinationCapacityGroupFits(4    allocation_mask: bits(4)) => boolean5begin6    var additional: integer = 0;7    for binding = 0 to PTO_BUNDLE_TILE_BINDING_COUNT - 1 do8        if _BundleTileBindings[[binding]].valid &&9           _BundleTileBindings[[binding]].destination_valid &&10           !_BundleTileBindings[[binding]].destination_allocated_by_bundle then11            additional = additional + BundleLocalDestinationAllocationBytes(12                binding as BundleTileBindingIndex);13        end;14    end;15    for pe = 0 to PTO_MODEL_MEMORY_AGENTS - 1 do16        if allocation_mask[PTOPEMaskBitOfPEIdentity(pe)] == '1' &&17           TileCapacityInUseForPE(pe) + additional >18               TileCapacityLimitBytes() then19            return FALSE;20        end;21    end;22    return TRUE;23end;24
25func ResolveBundleTMATMULCubeDestinationGroup(26    m: integer {1..65535}, n: integer {1..65535},27    accumulator_type: TileDataType,28    output_type: TileDataType,29    primary_layout: TileLayout,30    allocation_mask: bits(4)) => boolean31begin32    assert allocation_mask != Zeros{4};33    var reserved: array [[PTO_TILE_REGISTER_COUNT]] of boolean;34    var resolved: array [[PTO_BUNDLE_TILE_BINDING_COUNT]] of TileIndex;35    for index = 0 to PTO_TILE_REGISTER_COUNT - 1 do36        reserved[[index]] = FALSE;37    end;38    for binding = 0 to PTO_BUNDLE_TILE_BINDING_COUNT - 1 do39        resolved[[binding]] = 0;40        if _BundleTileBindings[[binding]].valid &&41           _BundleTileBindings[[binding]].destination_valid &&42           !_BundleTileBindings[[binding]].destination_allocated_by_bundle then43            let capacity_bytes = BundleTileDestinationSizeBytes(44                binding as BundleTileBindingIndex);45            let hand = UInt(46                _BundleTileBindings[[binding]].destination_hand);47            var found = FALSE;48            for offset = 0 to 15 do49                let raw_index: integer = hand * 16 + offset;50                if !found && !_Tiles[[raw_index]].allocated &&51                   !reserved[[raw_index]] then52                    resolved[[binding]] = raw_index as TileIndex;53                    reserved[[raw_index]] = TRUE;54                    found = TRUE;55                end;56            end;57            if !found then58                SetFault(Fault_TileAllocation, ReadTPC());59                return FALSE;60            end;61        end;62    end;63    if !BundleTMATMULDestinationCapacityGroupFits(allocation_mask) then64        SetFault(Fault_TileAllocation, ReadTPC());65        return FALSE;66    end;67
68    var destination_ordinal: integer {0..3} = 0;69    for binding = 0 to PTO_BUNDLE_TILE_BINDING_COUNT - 1 do70        if _BundleTileBindings[[binding]].valid &&71           _BundleTileBindings[[binding]].destination_valid &&72           !_BundleTileBindings[[binding]].destination_allocated_by_bundle then73            let capacity_bytes = BundleTileDestinationSizeBytes(74                binding as BundleTileBindingIndex);75            if destination_ordinal == 0 then76                if !TileMatrixMLayoutLegal(primary_layout, m) ||77                   !TileCubeDescriptorShapeLegal(78                       capacity_bytes, m, n,79                       output_type, primary_layout) then80                    SetFault(Fault_TileAllocation, ReadTPC());81                    return FALSE;82                end;83            else84                let row_auxiliary =85                    _BundleFixedPointAttributes.row_max_en &&86                    destination_ordinal == 1;87                let group_auxiliary =88                    _BundleFixedPointAttributes.group_max_en &&89                    ((!_BundleFixedPointAttributes.row_max_en &&90                      destination_ordinal == 1) ||91                     (_BundleFixedPointAttributes.row_max_en &&92                      destination_ordinal == 2));93                let auxiliary_columns = if row_auxiliary then 194                    else if group_auxiliary then95                        BundleGroupMaxColumns(n)96                    else n;97                let rows = DerivedTileRows(98                    capacity_bytes, auxiliary_columns, accumulator_type);99                if !TileDescriptorShapeLegal(100                       capacity_bytes, auxiliary_columns, m,101                       auxiliary_columns, accumulator_type) ||102                   rows * auxiliary_columns > PTO_MODEL_TILE_ELEMENTS then103                    SetFault(Fault_TileAllocation, ReadTPC());104                    return FALSE;105                end;106            end;107            destination_ordinal =108                (destination_ordinal + 1) as integer {0..3};109        end;110    end;111
112    destination_ordinal = 0;113    for binding = 0 to PTO_BUNDLE_TILE_BINDING_COUNT - 1 do114        if _BundleTileBindings[[binding]].valid &&115           _BundleTileBindings[[binding]].destination_valid &&116           !_BundleTileBindings[[binding]].destination_allocated_by_bundle then117            let capacity_bytes = BundleTileDestinationSizeBytes(118                binding as BundleTileBindingIndex);119            if destination_ordinal == 0 then120                let configured = ConfigureCubeTileForMask(121                    resolved[[binding]], capacity_bytes, m, n,122                    output_type, primary_layout,123                    TileLocation_Matrix, allocation_mask);124                assert configured;125            else126                let row_auxiliary =127                    _BundleFixedPointAttributes.row_max_en &&128                    destination_ordinal == 1;129                let group_auxiliary =130                    _BundleFixedPointAttributes.group_max_en &&131                    ((!_BundleFixedPointAttributes.row_max_en &&132                      destination_ordinal == 1) ||133                     (_BundleFixedPointAttributes.row_max_en &&134                      destination_ordinal == 2));135                let auxiliary_columns = if row_auxiliary then 1136                    else if group_auxiliary then137                        BundleGroupMaxColumns(n)138                    else n;139                ConfigureTileForMask(140                    resolved[[binding]], capacity_bytes,141                    m, auxiliary_columns, m, auxiliary_columns,142                    accumulator_type, TileLayout_RowMajor,143                    TileLocation_Any, allocation_mask);144            end;145            _BundleTileBindings[[binding]].destination =146                resolved[[binding]];147            _BundleTileBindings[[binding]].destination_allocated_by_bundle =148                TRUE;149            destination_ordinal =150                (destination_ordinal + 1) as integer {0..3};151        end;152    end;153    return TRUE;154end;155
156func ResolveBundleTMATMULDestination(157    m: integer {1..65535}, n: integer {1..65535},158    accumulator_type: TileDataType,159    cube_primary: boolean,160    primary_layout: TileLayout,161    allocation_mask: bits(4)) => boolean162begin163    var destination_binding: BundleTileBindingIndex = 0;164    var destination_seen = FALSE;165    for binding = 0 to PTO_BUNDLE_TILE_BINDING_COUNT - 1 do166        if !destination_seen &&167           _BundleTileBindings[[binding]].valid &&168           _BundleTileBindings[[binding]].destination_valid then169            destination_binding = binding as BundleTileBindingIndex;170            destination_seen = TRUE;171        end;172    end;173    if !destination_seen then174        SetFault(Fault_BundleControl, ReadTPC());175        return FALSE;176    end;177    let output_type = if178        UInt(_BundleFixedPointAttributes.pre_quant_mode) == 0 then179        accumulator_type180    else181        BundleFPATROutputType(182            _BundleFixedPointAttributes.pre_quant_mode);183    let capacity_bytes = BundleTileDestinationSizeBytes(destination_binding);184    if cube_primary then185        return ResolveBundleTMATMULCubeDestinationGroup(186            m, n, accumulator_type, output_type, primary_layout,187            allocation_mask);188    else189        let rows = DerivedTileRows(capacity_bytes, n, output_type);190        if !TileDescriptorShapeLegal(191               capacity_bytes, n, m, n, output_type) ||192           !IsNonzeroPowerOfTwo(rows) ||193           rows * n > PTO_MODEL_TILE_ELEMENTS then194            SetFault(Fault_TileAllocation, ReadTPC());195            return FALSE;196        end;197    end;198    // Reserve the complete D/RowMax/GroupMax destination group first.  Only199    // the primary D descriptor is then converted to persistent CUBE state;200    // auxiliary outputs remain ordinary Local Tiles.201    if !ResolveBundleTileDestinationsWithShape(TRUE, m, n, n) then202        return FALSE;203    end;204    return TRUE;205end;206
207func ResolveBundleTMATMULDestination(208    m: integer {1..65535}, n: integer {1..65535},209    accumulator_type: TileDataType,210    cube_primary: boolean,211    primary_layout: TileLayout) => boolean212begin213    var allocation_mask = Zeros{4};214    for binding = 0 to PTO_BUNDLE_TILE_BINDING_COUNT - 1 do215        if allocation_mask == Zeros{4} &&216           _BundleTileBindings[[binding]].valid &&217           _BundleTileBindings[[binding]].destination_valid then218            allocation_mask = _BundleTileBindings[[binding]].pe_mask;219        end;220    end;221    return ResolveBundleTMATMULDestination(222        m, n, accumulator_type, cube_primary, primary_layout,223        allocation_mask);224end;225
226func ResolveBundleTMATMULDestination(227    m: integer {1..65535}, n: integer {1..65535},228    accumulator_type: TileDataType) => boolean229begin230    var allocation_mask = Zeros{4};231    for binding = 0 to PTO_BUNDLE_TILE_BINDING_COUNT - 1 do232        if allocation_mask == Zeros{4} &&233           _BundleTileBindings[[binding]].valid &&234           _BundleTileBindings[[binding]].destination_valid then235            allocation_mask = _BundleTileBindings[[binding]].pe_mask;236        end;237    end;238    return ResolveBundleTMATMULDestination(239        m, n, accumulator_type, FALSE, TileLayout_RowMajor,240        allocation_mask);241end;242