PTO-TILE-MODEL-LEGALITY-REDUCTION-AND-EXPANSION

1// PTO-UNIT: {"id":"PTO-TILE-MODEL-LEGALITY-REDUCTION-AND-EXPANSION","surface":"tile","classification":["model","legality","reduction-and-expansion"],"depends_on":["PTO-TILE-MODEL-LEGALITY-OPERAND-SCHEMA"]}2
3readonly func TileOperandsLegal_ExecuteTileReduction(4    operation: TileReductionOperation,5    axis: TileAxis,6    destination: TileIndex,7    source: TileIndex) => boolean8begin9    if destination == source ||10       !TileDescriptorLegal(destination) ||11       !TileSourceContentsDefined(source) then12        return FALSE;13    end;14
15    let destination_tile = _Tiles[[destination]];16    let source_tile = _Tiles[[source]];17    let index_reduction =18        operation == TileReduction_ARGMIN ||19        operation == TileReduction_ARGMAX;20    let source_type_legal =21        if index_reduction then22            TileArgReductionSourceDataTypeSupported(source_tile.data_type)23        else24            TileVecArithmeticDataTypeSupported(source_tile.data_type);25    if destination_tile.storage_kind != TileStorage_Numeric ||26       source_tile.storage_kind != TileStorage_Numeric ||27       destination_tile.layout != TileLayout_RowMajor ||28       source_tile.layout != TileLayout_RowMajor ||29       source_tile.valid_rows == 0 ||30       source_tile.valid_columns == 0 ||31       !source_type_legal ||32       !TileSourceEncodingsValid(source) then33        return FALSE;34    end;35
36    if index_reduction then37        if destination_tile.data_type != TileDataType_S32 &&38           destination_tile.data_type != TileDataType_U32 then39            return FALSE;40        end;41    elsif destination_tile.data_type != source_tile.data_type then42        return FALSE;43    end;44
45    if axis == TileAxis_Row then46        return destination_tile.rows >= source_tile.valid_rows &&47               destination_tile.columns == 1 &&48               destination_tile.valid_rows == source_tile.valid_rows &&49               destination_tile.valid_columns == 1;50    end;51    return destination_tile.rows >= 1 &&52           destination_tile.columns == source_tile.columns &&53           destination_tile.valid_rows == 1 &&54           destination_tile.valid_columns ==55               source_tile.valid_columns;56end;57
58pure func TileExpandExpdifTypePairLegal(59    source_type: TileDataType,60    destination_type: TileDataType) => boolean61begin62    return (source_type == TileDataType_FP16 &&63            (destination_type == TileDataType_FP16 ||64             destination_type == TileDataType_FP32)) ||65           (source_type == TileDataType_BF16 &&66            (destination_type == TileDataType_BF16 ||67             destination_type == TileDataType_FP32)) ||68           (source_type == TileDataType_FP32 &&69            destination_type == TileDataType_FP32);70end;71
72readonly func TileOperandsLegal_ExecuteTileExpand(73    operation: TileExpandOperation,74    axis: TileAxis,75    destination: TileIndex,76    source: TileIndex,77    broadcast_source: TileIndex) => boolean78begin79    if !TileDescriptorLegal(destination) ||80       !TileDescriptorLegal(source) ||81       !TileDescriptorLegal(broadcast_source) then82        return FALSE;83    end;84
85    let destination_tile = _Tiles[[destination]];86    let source_tile = _Tiles[[source]];87    let broadcast_tile = _Tiles[[broadcast_source]];88    let expdif = operation == TileExpand_EXPDIF;89    let operation_type_legal =90        if operation == TileExpand_COPY then91            TileVecArithmeticDataTypeSupported(destination_tile.data_type)92        else if expdif then93            TileExpandExpdifTypePairLegal(94                source_tile.data_type,95                destination_tile.data_type)96        else97            TileVecArithmeticDataTypeSupported(destination_tile.data_type);98    if destination_tile.storage_kind != TileStorage_Numeric ||99       destination_tile.layout != TileLayout_RowMajor ||100       destination_tile.valid_rows == 0 ||101       destination_tile.valid_columns == 0 ||102       !operation_type_legal ||103       broadcast_tile.storage_kind != TileStorage_Numeric ||104       broadcast_tile.data_type !=105           (if expdif then source_tile.data_type106            else destination_tile.data_type) ||107       broadcast_tile.layout != TileLayout_RowMajor ||108       !TileSourceContentsDefined(broadcast_source) ||109       (operation != TileExpand_COPY &&110        !TileSourceEncodingsValid(broadcast_source)) then111        return FALSE;112    end;113
114    if operation == TileExpand_COPY then115        if source != broadcast_source then116            return FALSE;117        end;118    elsif (if expdif then !TileLogicalShapeMatch(destination, source)119           else !TileShapeAndTypeMatch(destination, source)) ||120          source_tile.storage_kind != TileStorage_Numeric ||121          source_tile.layout != TileLayout_RowMajor ||122          !TileSourceContentsDefined(source) ||123          (operation != TileExpand_COPY &&124           !TileSourceEncodingsValid(source)) then125        return FALSE;126    end;127
128    let broadcast_shape_legal =129        if axis == TileAxis_Row then130            broadcast_tile.valid_rows == destination_tile.valid_rows &&131            broadcast_tile.valid_columns == 1 &&132            broadcast_tile.columns == 1133        else134            broadcast_tile.valid_rows == 1 &&135            broadcast_tile.valid_columns ==136                destination_tile.valid_columns &&137            broadcast_tile.columns == destination_tile.columns;138    if !broadcast_shape_legal then139        return FALSE;140    end;141
142    if operation == TileExpand_DIV &&143       TileDataTypeIsInteger(destination_tile.data_type) then144        return TileBroadcastPayloadNonzero(145            axis,146            source,147            broadcast_source);148    end;149    return TRUE;150end;151