Defining Types, Operators, and Passes#

In this chapter, we walk through the Rust library of the package set up in Getting Started with Rust Extensions: a complex type, four operators, and a pass that lowers them to scalar arithmetic. The excerpts come from src/lib.rs, which also contains the validation and traversal helpers they use.

Defining a Type#

ComplexType has one field, elem, the scalar type shared by the real and imaginary parts:

/// A complex number whose real and imaginary parts share a scalar type.
#[rs_class("my_ext.ComplexType", final)]
pub struct ComplexTypeObj {
    base: TypeObj,
    pub elem: PrimType,
}
def_object_ref!(pub ComplexType, Type, ComplexTypeObj);

tvm_ffi::static_init_block! {
    register_type::<ComplexType>()?;
}

As in C++, the type is defined in two parts: ComplexTypeObj is the object that holds the fields, and ComplexType is a reference-counted reference to it, which IR nodes store and functions pass around. base: TypeObj makes the object an IR type, and rs_class registers its fields for reflection, through which Python and IR JSON loading construct the type as they construct a C++ one.

def_object_ref!(pub ComplexType, Type, ComplexTypeObj) defines the reference. Its arguments are:

pub ComplexType: the name and visibility of the reference.
Type: the reference to the parent object, TypeObj. A ComplexType converts into a Type with into() wherever a Type is expected, such as the result of type inference.
ComplexTypeObj: the object it refers to. A ComplexTypeObj converts into a ComplexType with into(), and a ComplexType dereferences to it, so ty.elem reads the field.

infer_complex in the next section uses both conversions: it builds a ComplexTypeObj, turns it into a ComplexType, and returns it as a Type.

The static_init_block! runs register_type when the library loads, as a TVM_FFI_STATIC_INIT_BLOCK does for C++, and Python then sees the type as my_ext.Complex.

A ComplexType records the scalar type of the two parts, and a complex(re, im) call supplies their values as its operands. elem can be any scalar float type, such as float32, float64 or bfloat16; this tutorial uses float32.

Defining Operators#

To support operations on the new type, we need to define new operators. In our case, complex arithmetic needs four: one builds a complex value from two scalars, one multiplies complex values, and two read the parts back as scalars so that results can be stored in ordinary float32 buffers. Like the type, they are registered under the my_ext namespace:

Operator	Meaning	Result type (`float32` parts)
`complex(re, im)`	Construct a complex value	`my_ext.Complex(T.float32)`
`cmul(z, w)`	Multiply two complex values	`my_ext.Complex(T.float32)`
`real(z)`	Extract the real part	`T.float32`
`imag(z)`	Extract the imaginary part	`T.float32`

The constructor is registered as an OpDef chain in a static_init_block!. Each step updates the operator registry, and the chain completes when its statement ends:

tvm_ffi::static_init_block! {
    OpDef::new("my_ext.complex", "Construct the complex number re + im * i.")
        .signature((
            sig::arg::<PrimExpr>("re", "The real part."),
            sig::arg::<PrimExpr>("im", "The imaginary part."),
        ))
        .set_infer_type(infer_complex)
        .set_validator(validate_call)
        .set_attr(CALL_EFFECT_KIND, CallEffectKind::Pure);
}

signature declares the arguments of the operator and of its Python and TVMScript builder. set_infer_type attaches a named Rust function that checks the arguments and computes the result type:

fn infer_complex(call: &CallObj) -> Result<Type> {
    check_call(call, 2)?;
    let elem = float_part(&call.args.get(0)?.ty)?;
    let im = float_part(&call.args.get(1)?.ty)?;
    require(elem.dtype == im.dtype, "my_ext.complex expects parts of the same type")?;
    let ty: ComplexType = ComplexTypeObj::new(TypeObj::new(call.loc.clone()), elem).into();
    Ok(ty.into())
}

set_validator repeats these checks for a call whose result type is given explicitly rather than inferred. Pure marks the operator as computing a value without writing memory.

infer_complex requires both parts to be scalar floats of the same type, and converts neither. The other three operators follow the same pattern. cmul requires two complex arguments of the same type and returns it; real and imag require one and return its elem. Passing an ordinary scalar to cmul therefore fails with TypeError: my_ext.cmul expects a complex value.

Writing a Pass#

Although the new type and operators can be used like built-in ones when writing programs, the existing passes know nothing about them: they can neither optimize them further nor lower them to what the backend supports. We therefore add a pass of our own to handle them. In our case, LowerComplex applies three rewrite rules:

cmul(complex(a, b), complex(c, d))  ->  complex(a*c - b*d, a*d + b*c)
real(complex(a, b))                 ->  a
imag(complex(a, b))                 ->  b

A complex local variable z = complex(a, b) is split into two scalar variables, z_re = a and z_im = b, and later uses of z resolve to that pair. This lets the rules apply to the step-by-step function from Getting Started with Rust Extensions, whose loop body becomes:

z_re: T.let[T.float32] = A[i]
z_im: T.let[T.float32] = B[i]
w_re: T.let[T.float32] = C[i]
w_im: T.let[T.float32] = D[i]
product_re: T.let[T.float32] = z_re * w_re - z_im * w_im
product_im: T.let[T.float32] = z_re * w_im + z_im * w_re
R[i] = product_re
Im[i] = product_im

Only scalar arithmetic remains. The pass looks up the four operators by name, as C++ does with Op::Get:

struct ComplexOps {
    complex: Op,
    cmul: Op,
    real: Op,
    imag: Op,
}

impl ComplexOps {
    fn get() -> Result<Self> {
        Ok(Self {
            complex: Op::get("my_ext.complex")?,
            cmul: Op::get("my_ext.cmul")?,
            real: Op::get("my_ext.real")?,
            imag: Op::get("my_ext.imag")?,
        })
    }

    /// Whether `op`, such as a call's `op`, is one of the four operators.
    fn contains(&self, op: &Expr) -> bool {
        [&self.complex, &self.cmul, &self.real, &self.imag].into_iter().any(|own| op.same_as(own))
    }
}

The operators are registered when the library loads, and Op::get returns the registry’s single object for a name. call.op.same_as(&ops.cmul) therefore tests whether a call’s op is my_ext.cmul by comparing pointers. A lookup goes through the registry, so the pass fetches the operators once per run and keeps them. Lowering implements the rewrites with them:

struct Lowering {
    ops: ComplexOps,
    bindings: Vec<(Var, Var, Var)>,
}
impl Lowering {
    fn parts(&self, expr: &Expr) -> Result<(Expr, Expr)> { ... }
}
#[dispatch(map)]
impl Lowering {
    fn map_sequence(&mut self, seq: SeqStmt) -> Result<Any> { ... }
    fn map_call(&mut self, call: Call) -> Result<Any> {
        let ops = &self.ops;
        if ops.contains(&call.op) {
            validate_call(&call)?;
        }
        if call.op.same_as(&ops.cmul) {
            let (a, b) = self.parts(&call.args.get(0)?)?;
            let (c, d) = self.parts(&call.args.get(1)?)?;
            // (a + bi)(c + di) = (ac - bd) + (ad + bc)i
            let re = &a * &c - &b * &d;
            let im = &a * &d + &b * &c;
            return Ok(Call::new(
                None,
                ops.complex.clone(),
                [re, im],
                None,
                Vec::<Type>::new(),
                call.loc.clone(),
            )?
            .into());
        }
        if call.op.same_as(&ops.real) || call.op.same_as(&ops.imag) {
            let (re, im) = self.parts(&call.args.get(0)?)?;
            return Ok(if call.op.same_as(&ops.real) { re.into() } else { im.into() });
        }
        Ok(call.into())
    }

    fn map_bind(&mut self, bind: Bind) -> Result<Any> {
        if let Some(ty) = bind.var.ty.downcast_ref::<ComplexTypeObj>() {
            let (re, im) = self.parts(&bind.value)?;
            let name = bind.var.name.as_str();
            let re_var = Var::new(&format!("{name}_re"), ty.elem.clone(), bind.var.loc.clone());
            let im_var = Var::new(&format!("{name}_im"), ty.elem.clone(), bind.var.loc.clone());
            self.bindings.push((bind.var.clone(), re_var.clone(), im_var.clone()));
            let binds = [
                Bind::new(re_var, re, bind.loc.clone())?,
                Bind::new(im_var, im, bind.loc.clone())?,
            ];
            return Ok(SeqStmt::new(binds, bind.loc.clone()).into());
        }
        Ok(bind.into())
    }
}

ops holds the four operators, bindings records each complex variable with its two scalar replacements, and parts reads the real and imaginary parts of an operand, either a complex(re, im) call or a variable already split. #[dispatch(map)] sends each node to the map_ method of its type, and leaves other nodes unchanged: map_call applies the three rules, map_bind splits a complex binding, and map_sequence flattens the sequence that a split binding leaves inside another. The value-only callback passes each node to lowering.dispatch and returns Unchanged when no handler matches; these rewrites do not need a definition-region argument.

The pass is built from an ordinary named function:

fn lower_function(
    func: Function,
    _module: IRModule,
    _ctx: PassContext,
) -> Result<Option<Function>> {
    let mut lowering = Lowering { ops: ComplexOps::get()?, bindings: Vec::new() };
    let rewritten = tvm_ffi::structural_map(
        func,
        |value: AnyView<'_>| lowering.dispatch(value).unwrap_or_else(|| Ok(Unchanged.into())),
        WalkOrder::PostOrder,
    )?;
    let ops = &lowering.ops;
    let checked = tvm_ffi::structural_map(
        rewritten,
        (check_type, |call| check_call_lowered(ops, call), check_expr),
        WalkOrder::PostOrder,
    )?;
    let function: Function = checked.try_into()?;
    Ok(Some(function))
}

It receives a function, its module, and the pass context, and returns the rewritten function; Some replaces the original. The first structural_map walks the function with Lowering. PostOrder visits children before their parent, so the operands of cmul are lowered before cmul itself. The second structural_map only checks the result. It takes a tuple of three functions, and calls each on the nodes of its parameter type; the closure passes the operators to check_call_lowered:

/// Reject a complex type, wherever it appears: on a variable, a parameter or an expression.
fn check_type(ty: Type) -> Result<Type> { ... }

/// Reject an expression whose type is complex.
fn check_expr(expr: Expr) -> Result<Expr> { ... }

/// Reject a call to one of the four complex operators that the rewrite did not consume.
fn check_call_lowered(ops: &ComplexOps, call: Call) -> Result<Call> { ... }

Each returns its node unchanged and raises a ValueError on an unsupported use, such as a complex function parameter, so the pass reports such a use before code generation.

Finally, FunctionPass::new wraps the function as a pass, and the factory is exported so that Python can call it as lib.lower_complex():

fn lower_complex() -> Result<Pass> {
    let pass = FunctionPass::new(lower_function, PassInfo::new(0, "my_ext.LowerComplex"));
    Ok(Pass::from(&pass))
}

tvm_ffi::tvm_ffi_dll_export_typed_func!(lower_complex, lower_complex);

Next, Using and Compiling the Extension loads the library and compiles the function.

Defining Types, Operators, and Passes

Contents

Defining Types, Operators, and Passes#

Defining a Type#

Defining Operators#

Writing a Pass#