Function test_example_npu_runtime_execution

tests/python/contrib/test_example_npu.py:225–260

Test end-to-end execution with the example NPU runtime

()

@example_npu_enabled
def test_example_npu_runtime_execution():
    """Test end-to-end execution with the example NPU runtime"""
    import tvm.relax.backend.contrib.example_npu

    # Create simple test inputs
    np.random.seed(42)
    x_np = np.random.randn(2, 4).astype("float32")
    w_np = np.random.randn(4, 8).astype("float32")

    # Expected output (computed with NumPy)
    expected = np.maximum(0, np.matmul(x_np, w_np))

    # Build and run with example NPU backend
    mod = MatmulReLU
    patterns = get_patterns_with_prefix("example_npu")

    # Apply transformations
    mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True)(mod)
    mod = MergeCompositeFunctions()(mod)
    mod = RunCodegen()(mod)

    # Build the module
    target = tvm.target.Target("llvm")
    with tvm.transform.PassContext(opt_level=3):
        built = relax.build(mod, target)

    # Create VM and run
    vm = relax.VirtualMachine(built, tvm.cpu())

    x_tvm = tvm.runtime.tensor(x_np, tvm.cpu())
    w_tvm = tvm.runtime.tensor(w_np, tvm.cpu())

    result = vm["main"](x_tvm, w_tvm)

    # Verify the result shape is correct (the runtime is a stub and does not compute numerically)
    assert result.numpy().shape == expected.shape


if __name__ == "__main__":