PTO-TILE-MODEL-EXECUTION-UNARY

1// PTO-UNIT: {"id":"PTO-TILE-MODEL-EXECUTION-UNARY","surface":"tile","classification":["model","execution","unary"],"depends_on":["PTO-TILE-MODEL-EXECUTION-ELEMENTWISE"]}2
3pure func TileIntegerUnaryValue(4    operation: TileUnaryOperation,5    data_type: TileDataType,6    value: Word) => Word7begin8    assert TileDataTypeIsInteger(data_type);9
10    let unsigned_value = TileUnsignedElementValue(value, data_type);11    let signed_value = TileIntegerOperandValue(value, data_type);12    case operation of13        when TileUnary_ABS =>14            if TileDataTypeIsSigned(data_type) &&15               SInt(signed_value) < 0 then16                return TileUnsignedElementValue(17                    Zeros{PTO_XLEN} - signed_value,18                    data_type);19            else20                return unsigned_value;21            end;22        when TileUnary_NOT =>23            return TileUnsignedElementValue(24                NOT unsigned_value,25                data_type);26        when TileUnary_NEG =>27            return TileUnsignedElementValue(28                Zeros{PTO_XLEN} - unsigned_value,29                data_type);30        when TileUnary_RELU =>31            if TileDataTypeIsSigned(data_type) &&32               SInt(signed_value) < 0 then33                return Zeros{PTO_XLEN};34            else35                return unsigned_value;36            end;37        otherwise =>38            unreachable;39    end;40end;41
42pure func TileUnaryUsesClosedElementwiseContract(43    operation: TileUnaryOperation) => boolean44begin45    return operation == TileUnary_ABS ||46           operation == TileUnary_NOT ||47           operation == TileUnary_NEG ||48           operation == TileUnary_RELU;49end;50
51pure func TileUnaryUsesSFUElementwiseContract(52    operation: TileUnaryOperation) => boolean53begin54    return operation == TileUnary_EXP ||55           operation == TileUnary_LOG ||56           operation == TileUnary_RECIP ||57           operation == TileUnary_SQRT ||58           operation == TileUnary_RSQRT;59end;60
61pure func TileUnaryUsesCompleteElementwiseSchema(62    operation: TileUnaryOperation) => boolean63begin64    return TileUnaryUsesClosedElementwiseContract(operation) ||65           TileUnaryUsesSFUElementwiseContract(operation);66end;67
68pure func TileUnaryScalarIntegerDataType(69    data_type: TileDataType) => boolean70begin71    return data_type == TileDataType_S64 ||72           data_type == TileDataType_S32 ||73           data_type == TileDataType_S16 ||74           data_type == TileDataType_S8 ||75           data_type == TileDataType_U64 ||76           data_type == TileDataType_U32 ||77           data_type == TileDataType_U16 ||78           data_type == TileDataType_U8;79end;80
81pure func TileUnaryDataTypeSupported(82    operation: TileUnaryOperation,83    data_type: TileDataType) => boolean84begin85    if operation == TileUnary_NOT then86        return TileVecScalarIntegerDataTypeSupported(data_type);87    end;88    if operation == TileUnary_ABS then89        return TileVecArithmeticDataTypeSupported(data_type);90    end;91    if operation == TileUnary_NEG then92        return TileTNegDataTypeSupported(data_type);93    end;94    if operation == TileUnary_RELU then95        return TileTReluDataTypeSupported(data_type);96    end;97    if TileUnaryUsesSFUElementwiseContract(operation) then98        return TileFloatingElementwiseDataTypeSupported(data_type);99    end;100    return FALSE;101end;102
103pure func TileFloatingUnaryCarrier(104    data_type: TileDataType,105    value: Word) => Word106begin107    case data_type of108        when TileDataType_FP64 =>109            return value;110        when TileDataType_FP32, TileDataType_TF32, TileDataType_HF32 =>111            return ZeroExtend{PTO_XLEN}(value[31:0]);112        when TileDataType_FP16, TileDataType_BF16 =>113            return ZeroExtend{PTO_XLEN}(value[15:0]);114        when TileDataType_E4M3, TileDataType_E5M2 =>115            return ZeroExtend{PTO_XLEN}(value[7:0]);116        otherwise =>117            unreachable;118    end;119end;120
121pure func TileFloatingUnarySignMask(122    data_type: TileDataType) => Word123begin124    case data_type of125        when TileDataType_FP64 =>126            return Zeros{PTO_XLEN} + 0x8000000000000000;127        when TileDataType_FP32, TileDataType_TF32, TileDataType_HF32 =>128            return Zeros{PTO_XLEN} + 0x80000000;129        when TileDataType_FP16, TileDataType_BF16 =>130            return Zeros{PTO_XLEN} + 0x8000;131        when TileDataType_E4M3, TileDataType_E5M2 =>132            return Zeros{PTO_XLEN} + 0x80;133        otherwise =>134            unreachable;135    end;136end;137
138pure func TileFloatingUnaryOneEncoding(139    data_type: TileDataType) => Word140begin141    case data_type of142        when TileDataType_FP64 =>143            return Zeros{PTO_XLEN} + 0x3ff0000000000000;144        when TileDataType_FP32, TileDataType_TF32,145             TileDataType_HF32 =>146            return Zeros{PTO_XLEN} + 0x3f800000;147        when TileDataType_FP16 =>148            return Zeros{PTO_XLEN} + 0x3c00;149        when TileDataType_BF16 =>150            return Zeros{PTO_XLEN} + 0x3f80;151        when TileDataType_E4M3 =>152            return Zeros{PTO_XLEN} + 0x38;153        when TileDataType_E5M2 =>154            return Zeros{PTO_XLEN} + 0x3c;155        otherwise =>156            unreachable;157    end;158end;159
160pure func TileFloatingUnaryInfinityEncoding(161    data_type: TileDataType,162    negative: boolean) => (boolean, Word)163begin164    var positive = Zeros{PTO_XLEN};165    case data_type of166        when TileDataType_FP64 =>167            positive = Zeros{PTO_XLEN} + 0x7ff0000000000000;168        when TileDataType_FP32, TileDataType_TF32,169             TileDataType_HF32 =>170            positive = Zeros{PTO_XLEN} + 0x7f800000;171        when TileDataType_FP16 =>172            positive = Zeros{PTO_XLEN} + 0x7c00;173        when TileDataType_BF16 =>174            positive = Zeros{PTO_XLEN} + 0x7f80;175        when TileDataType_E5M2 =>176            positive = Zeros{PTO_XLEN} + 0x7c;177        when TileDataType_E4M3 =>178            return (FALSE, Zeros{PTO_XLEN});179        otherwise =>180            unreachable;181    end;182    if negative then183        return (184            TRUE,185            positive OR TileFloatingUnarySignMask(data_type));186    end;187    return (TRUE, positive);188end;189
190pure func TileFloatingUnaryUnboundedResult(191    data_type: TileDataType,192    negative: boolean) => Word193begin194    let (infinity_available, infinity) =195        TileFloatingUnaryInfinityEncoding(196            data_type,197            negative);198    if infinity_available then199        return infinity;200    end;201
202    // E4M3 has no infinity encoding.  These SFU operations do not admit203    // saturation, so the OFP8 non-saturating destination result is NaN.204    let (nan_available, quiet_nan) =205        HardwareNumericCanonicalNaNResult(data_type);206    assert nan_available;207    return quiet_nan;208end;209
210pure func TileNumericValueClassIsNegative(211    value_class: NumericValueClass) => boolean212begin213    return value_class == NumericValue_NegativeZero ||214           value_class == NumericValue_NegativeSubnormal ||215           value_class == NumericValue_NegativeNormal ||216           value_class == NumericValue_NegativeInfinity;217end;218
219pure func TileSFUUnarySpecialValue(220    operation: TileUnaryOperation,221    data_type: TileDataType,222    value: Word) => (boolean, Word, bits(5))223begin224    assert TileUnaryUsesSFUElementwiseContract(operation);225    assert TileUnaryDataTypeSupported(operation, data_type);226    assert TileNumericEncodingValid(data_type, value);227
228    let carrier = TileFloatingUnaryCarrier(data_type, value);229    let value_class = TileNumericValueClass(data_type, carrier);230    let (nan_available, quiet_nan) =231        HardwareNumericCanonicalNaNResult(data_type);232    let (zero_available, positive_zero, negative_zero) =233        HardwareNumericSignedZeroEncodings(data_type);234    assert nan_available;235    assert zero_available;236
237    if NumericValueClassIsNaN(value_class) then238        var nan_flags = Zeros{5};239        if value_class == NumericValue_SignalingNaN then240            nan_flags = Zeros{5} + 1;241        end;242        return (TRUE, quiet_nan, nan_flags);243    end;244
245    if operation == TileUnary_EXP then246        if NumericValueClassIsZero(value_class) then247            return (248                TRUE,249                TileFloatingUnaryOneEncoding(data_type),250                Zeros{5});251        elsif value_class == NumericValue_PositiveInfinity then252            return (TRUE, carrier, Zeros{5});253        elsif value_class == NumericValue_NegativeInfinity then254            return (TRUE, positive_zero, Zeros{5});255        end;256    elsif operation == TileUnary_LOG then257        if carrier == TileFloatingUnaryOneEncoding(data_type) then258            return (TRUE, positive_zero, Zeros{5});259        elsif NumericValueClassIsZero(value_class) then260            return (261                TRUE,262                TileFloatingUnaryUnboundedResult(data_type, TRUE),263                Zeros{5} + 2);264        elsif value_class == NumericValue_PositiveInfinity then265            return (TRUE, carrier, Zeros{5});266        elsif TileNumericValueClassIsNegative(value_class) then267            return (TRUE, quiet_nan, Zeros{5} + 1);268        end;269    elsif operation == TileUnary_RECIP then270        if NumericValueClassIsZero(value_class) then271            let negative = value_class == NumericValue_NegativeZero;272            return (273                TRUE,274                TileFloatingUnaryUnboundedResult(data_type, negative),275                Zeros{5} + 2);276        elsif NumericValueClassIsInfinity(value_class) then277            if value_class == NumericValue_NegativeInfinity then278                return (TRUE, negative_zero, Zeros{5});279            end;280            return (TRUE, positive_zero, Zeros{5});281        end;282    elsif operation == TileUnary_SQRT then283        if NumericValueClassIsZero(value_class) ||284           value_class == NumericValue_PositiveInfinity then285            return (TRUE, carrier, Zeros{5});286        elsif TileNumericValueClassIsNegative(value_class) then287            return (TRUE, quiet_nan, Zeros{5} + 1);288        end;289    else290        assert operation == TileUnary_RSQRT;291        if NumericValueClassIsZero(value_class) then292            let negative = value_class == NumericValue_NegativeZero;293            return (294                TRUE,295                TileFloatingUnaryUnboundedResult(data_type, negative),296                Zeros{5} + 2);297        elsif value_class == NumericValue_PositiveInfinity then298            return (TRUE, positive_zero, Zeros{5});299        elsif TileNumericValueClassIsNegative(value_class) then300            return (TRUE, quiet_nan, Zeros{5} + 1);301        end;302    end;303    return (FALSE, Zeros{PTO_XLEN}, Zeros{5});304end;305
306pure func TileFloatingUnaryValue(307    operation: TileUnaryOperation,308    data_type: TileDataType,309    value: Word) => (Word, boolean)310begin311    assert TileDataTypeIsFloating(data_type);312    assert TileNumericEncodingValid(data_type, value);313
314    let carrier = TileFloatingUnaryCarrier(data_type, value);315    let sign_mask = TileFloatingUnarySignMask(data_type);316    if operation == TileUnary_ABS then317        return (carrier AND NOT sign_mask, FALSE);318    elsif operation == TileUnary_NEG then319        return (carrier XOR sign_mask, FALSE);320    end;321
322    assert operation == TileUnary_RELU;323    let value_class = TileNumericValueClass(data_type, carrier);324    case value_class of325        when NumericValue_NegativeZero,326             NumericValue_NegativeSubnormal,327             NumericValue_NegativeNormal,328             NumericValue_NegativeInfinity,329             NumericValue_PositiveZero =>330            return (Zeros{PTO_XLEN}, FALSE);331        when NumericValue_PositiveSubnormal,332             NumericValue_PositiveNormal,333             NumericValue_PositiveInfinity =>334            return (carrier, FALSE);335        when NumericValue_QuietNaN, NumericValue_SignalingNaN =>336            let (available, quiet_nan) =337                HardwareNumericCanonicalNaNResult(data_type);338            assert available;339            return (340                quiet_nan,341                value_class == NumericValue_SignalingNaN);342        otherwise =>343            unreachable;344    end;345end;346
347pure func TileFixedUnaryValue(348    operation: TileUnaryOperation,349    data_type: TileDataType,350    value: Word) => (Word, boolean)351begin352    assert TileUnaryUsesClosedElementwiseContract(operation);353    assert TileUnaryDataTypeSupported(operation, data_type);354    if TileDataTypeIsInteger(data_type) then355        return (356            TileIntegerUnaryValue(operation, data_type, value),357            FALSE);358    end;359    return TileFloatingUnaryValue(operation, data_type, value);360end;361
362func TileUnaryValue(operation: TileUnaryOperation, value: Word) => Word363begin364    case operation of365        when TileUnary_ABS =>366            if SInt(value) < 0 then367                return Zeros{PTO_XLEN} - value;368            else369                return value;370            end;371        when TileUnary_NOT =>372            return NOT value;373        when TileUnary_NEG =>374            return Zeros{PTO_XLEN} - value;375        when TileUnary_RELU =>376            if SInt(value) < 0 then377                return Zeros{PTO_XLEN};378            else379                return value;380            end;381        when TileUnary_SQRT =>382            return TileSquareRoot(value);383        when TileUnary_LOG =>384            return TileLogarithm(value);385        when TileUnary_RECIP =>386            return TileReciprocal(value);387        when TileUnary_EXP =>388            return TileExponential(value);389        when TileUnary_RSQRT =>390            return TileReciprocalSquareRoot(value);391    end;392end;393
394impdef func TileProfileUnary(395    op: TileUnaryOperation,396    data_type: TileDataType,397    value: Word) => (Word, bits(5))398begin399    if TileUnaryUsesClosedElementwiseContract(op) then400        let (result, invalid) = TileFixedUnaryValue(401            op,402            data_type,403            value);404        return (405            result,406            if invalid then Zeros{5} + 1 else Zeros{5});407    end;408    return (409        TileUnaryValue(op, value),410        Zeros{5});411end;412
413func ExecuteTileUnary(414    operation: TileUnaryOperation,415    destination: TileIndex,416    source: TileIndex)417begin418    let source_tile = _Tiles[[source]];419    assert source_tile.allocated;420    assert TileShapesMatch(_Tiles[[destination]], source_tile);421    assert _Tiles[[destination]].data_type == source_tile.data_type;422
423    var result_tile = _Tiles[[destination]];424    var flags = Zeros{5};425    for row = 0 to source_tile.valid_rows - 1 looplimit 65536 do426        for column = 0 to source_tile.valid_columns - 1 looplimit 65536 do427            let element = TileLogicalLinearIndex(428                source_tile,429                row as integer {0..65535},430                column as integer {0..65535});431            if TileUnaryUsesClosedElementwiseContract(operation) then432                let (result, element_invalid) = TileFixedUnaryValue(433                    operation,434                    source_tile.data_type,435                    TileReadLogicalElement(source_tile, element));436                result_tile = TileInfoWithLogicalElement(437                    result_tile, element, result);438                if element_invalid then439                    flags = flags OR (Zeros{5} + 1);440                end;441            else442                let (handled, special_result, special_flags) =443                    TileSFUUnarySpecialValue(444                        operation,445                        source_tile.data_type,446                        TileReadLogicalElement(source_tile, element));447                var result = special_result;448                var element_flags = special_flags;449                if !handled then450                    let (profile_result, profile_flags) = TileProfileUnary(451                        operation,452                        source_tile.data_type,453                        TileReadLogicalElement(source_tile, element));454                    result = profile_result;455                    element_flags = profile_flags;456                end;457                result_tile = TileInfoWithLogicalElement(458                    result_tile, element, result);459                flags = flags OR element_flags;460            end;461        end;462    end;463    _Tiles[[destination]] = result_tile;464    MarkTileValidRegionDefined(destination);465    if TileUnaryUsesCompleteElementwiseSchema(operation) then466        ApplyTilePadding(destination, CurrentBundlePadValue());467    end;468    ScalarFPRecordFlags(flags);469end;470