Using and Compiling the Extension

Using and Compiling the Extension#

In this chapter, we look at the Python package that loads the library defined in Defining Types, Operators, and Passes, then use it to compile the complex multiplication function and check the result against NumPy.

The Python Package#

In Defining Types, Operators, and Passes, the type and operators were registered in static init blocks, which run as soon as the library loads. To write them in TVMScript, Python also needs a module that holds their builders. Definitions under tirx. go to T, but an extension namespace such as my_ext is published only into a module that the extension declares. The package’s __init__.py therefore declares itself as that module before it loads the library:

_init_exposed_op_api("my_ext", __name__, wrapper=auto_emit_void_call)
_init_exposed_type_api("my_ext", __name__)
_LIB = load_lib_module("my-ext", "my_ext", extra_lib_paths=[Path(__file__).parent])

The two _init_exposed_*_api calls, internal to tvm_tirx.ir for now, publish the operators and the type into my_ext; auto_emit_void_call gives its builders the TVMScript behavior of the T builders. load_lib_module finds the library that the package installs next to __init__.py. To parse printed IR that uses the extension, pass the package as well, as in extra_vars={"T": T, "my_ext": my_ext}.

The registered type and operators let the IR express complex values; LowerComplex turns them into a form the backend can compile. The default pipeline does not know about the pass, so the package registers a pipeline, also named my_ext, that runs it first:

def _pipeline(*, target=None):
    lower, finalize_host, finalize_device = tirx.get_tir_pipeline("default", target=target)
    return ir.transform.Sequential([lower_complex(), lower]), finalize_host, finalize_device


tirx.register_tir_pipeline("my_ext", _pipeline)

lower_complex returns the pass that the library exports.

Compile the Function#

run_example.py imports my_ext, which loads the library and registers the pipeline, and compiles the function from Getting Started with Rust Extensions with it:

compiled = tvm_tirx.compile(complex_product, target="llvm", pipeline="my_ext")

Compiling with the default pipeline instead fails in LLVM code generation, which has no complex type.

Check the Result#

A successful compilation shows that LowerComplex left no complex values for the backend. To check that its rewrite also computes the right numbers, the script runs the compiled function on random float32 inputs and compares both outputs with NumPy’s complex multiplication:

rng = np.random.default_rng(0)
a, b, c, d = rng.standard_normal((4, N)).astype("float32")
inputs = [runtime.tensor(x) for x in (a, b, c, d)]
real_out = runtime.empty((N,), "float32")
imag_out = runtime.empty((N,), "float32")
compiled(*inputs, real_out, imag_out)

expected = (a + 1j * b) * (c + 1j * d)
np.testing.assert_allclose(real_out.numpy(), expected.real, rtol=1e-5, atol=1e-6)
np.testing.assert_allclose(imag_out.numpy(), expected.imag, rtol=1e-5, atol=1e-6)

With the package installed as in Getting Started with Rust Extensions, run the script:

python examples/extension/run_example.py

A successful run prints:

Complex multiplication matches NumPy.

tests/python/ir/test_op_rust.py runs the script too, and also checks printing and parsing, invalid calls, chained products, and that lowering leaves no complex types or operators.

Next Steps#

Complex addition or conjugation needs one more operator and one more rewrite rule in LowerComplex, registered and written the same way. The example keeps complex values inside a function: parameters and outputs are ordinary arrays, and complex numbers have no runtime storage format or calling convention. See the Rust API reference for the full Rust API.