Convert arith to riscv snitch

convert_arith_to_riscv_snitch

lower_arith_addf = LowerBinaryFloatVectorOp(arith.AddfOp, riscv.FAddDOp, riscv_snitch.VFAddSOp, riscv_snitch.VFAddHOp) module-attribute

lower_arith_subf = LowerBinaryFloatVectorOp(arith.SubfOp, riscv.FSubDOp, riscv_snitch.VFSubSOp, riscv_snitch.VFSubHOp) module-attribute

lower_arith_mulf = LowerBinaryFloatVectorOp(arith.MulfOp, riscv.FMulDOp, riscv_snitch.VFMulSOp, riscv_snitch.VFMulHOp) module-attribute

LowerBinaryFloatVectorOp dataclass

Bases: RewritePattern

Source code in xdsl/backend/riscv/lowering/convert_arith_to_riscv_snitch.py

@dataclass
class LowerBinaryFloatVectorOp(RewritePattern):
    arith_op_cls: type[arith.FloatingPointLikeBinaryOperation]
    riscv_d_op_cls: type[riscv.RdRsRsFloatOperationWithFastMath]
    riscv_snitch_v_f32_op_cls: type[riscv.RdRsRsFloatOperationWithFastMath]
    riscv_snitch_v_f16_op_cls: type[riscv.RdRsRsFloatOperationWithFastMath]

    def match_and_rewrite(self, op: Operation, rewriter: PatternRewriter) -> None:
        if not isinstance(op, self.arith_op_cls):
            return

        operand_type = op.result.type
        if not isinstance(operand_type, VectorType):
            return
        shape = operand_type.shape
        count = prod(dim.data for dim in shape.data)

        operand_type = cast(VectorType[Any], operand_type)
        scalar_type = operand_type.element_type

        lhs = UnrealizedConversionCastOp.get(
            (op.lhs,), (riscv.Registers.UNALLOCATED_FLOAT,)
        )
        rhs = UnrealizedConversionCastOp.get(
            (op.rhs,), (riscv.Registers.UNALLOCATED_FLOAT,)
        )

        match scalar_type:
            case Float64Type():
                if count != 1:
                    return
                cls = self.riscv_d_op_cls
            case Float32Type():
                if count != 2:
                    return
                cls = self.riscv_snitch_v_f32_op_cls
            case Float16Type():
                if count != 4:
                    return
                cls = self.riscv_snitch_v_f16_op_cls
            case _:
                raise ValueError(f"Unexpected float type {op.lhs.type}")

        rv_flags = riscv.FastMathFlagsAttr(op.fastmath.data)

        new_op = cls(lhs, rhs, rd=riscv.Registers.UNALLOCATED_FLOAT, fastmath=rv_flags)
        cast_op = UnrealizedConversionCastOp.get((new_op.rd,), (op.result.type,))

        rewriter.replace_op(op, (lhs, rhs, new_op, cast_op))

arith_op_cls: type[arith.FloatingPointLikeBinaryOperation] instance-attribute

riscv_d_op_cls: type[riscv.RdRsRsFloatOperationWithFastMath] instance-attribute

riscv_snitch_v_f32_op_cls: type[riscv.RdRsRsFloatOperationWithFastMath] instance-attribute

riscv_snitch_v_f16_op_cls: type[riscv.RdRsRsFloatOperationWithFastMath] instance-attribute

__init__(arith_op_cls: type[arith.FloatingPointLikeBinaryOperation], riscv_d_op_cls: type[riscv.RdRsRsFloatOperationWithFastMath], riscv_snitch_v_f32_op_cls: type[riscv.RdRsRsFloatOperationWithFastMath], riscv_snitch_v_f16_op_cls: type[riscv.RdRsRsFloatOperationWithFastMath]) -> None

match_and_rewrite(op: Operation, rewriter: PatternRewriter) -> None

Source code in xdsl/backend/riscv/lowering/convert_arith_to_riscv_snitch.py

def match_and_rewrite(self, op: Operation, rewriter: PatternRewriter) -> None:
    if not isinstance(op, self.arith_op_cls):
        return

    operand_type = op.result.type
    if not isinstance(operand_type, VectorType):
        return
    shape = operand_type.shape
    count = prod(dim.data for dim in shape.data)

    operand_type = cast(VectorType[Any], operand_type)
    scalar_type = operand_type.element_type

    lhs = UnrealizedConversionCastOp.get(
        (op.lhs,), (riscv.Registers.UNALLOCATED_FLOAT,)
    )
    rhs = UnrealizedConversionCastOp.get(
        (op.rhs,), (riscv.Registers.UNALLOCATED_FLOAT,)
    )

    match scalar_type:
        case Float64Type():
            if count != 1:
                return
            cls = self.riscv_d_op_cls
        case Float32Type():
            if count != 2:
                return
            cls = self.riscv_snitch_v_f32_op_cls
        case Float16Type():
            if count != 4:
                return
            cls = self.riscv_snitch_v_f16_op_cls
        case _:
            raise ValueError(f"Unexpected float type {op.lhs.type}")

    rv_flags = riscv.FastMathFlagsAttr(op.fastmath.data)

    new_op = cls(lhs, rhs, rd=riscv.Registers.UNALLOCATED_FLOAT, fastmath=rv_flags)
    cast_op = UnrealizedConversionCastOp.get((new_op.rd,), (op.result.type,))

    rewriter.replace_op(op, (lhs, rhs, new_op, cast_op))

ConvertArithToRiscvSnitchPass dataclass

Bases: ModulePass

Source code in xdsl/backend/riscv/lowering/convert_arith_to_riscv_snitch.py

class ConvertArithToRiscvSnitchPass(ModulePass):
    name = "convert-arith-to-riscv-snitch"

    def apply(self, ctx: Context, op: ModuleOp) -> None:
        walker = PatternRewriteWalker(
            GreedyRewritePatternApplier(
                [
                    lower_arith_addf,
                    lower_arith_subf,
                    lower_arith_mulf,
                ],
            ),
            apply_recursively=False,
        )
        walker.rewrite_module(op)

name = 'convert-arith-to-riscv-snitch' class-attribute instance-attribute

apply(ctx: Context, op: ModuleOp) -> None

Source code in xdsl/backend/riscv/lowering/convert_arith_to_riscv_snitch.py

def apply(self, ctx: Context, op: ModuleOp) -> None:
    walker = PatternRewriteWalker(
        GreedyRewritePatternApplier(
            [
                lower_arith_addf,
                lower_arith_subf,
                lower_arith_mulf,
            ],
        ),
        apply_recursively=False,
    )
    walker.rewrite_module(op)