Class DebugModule

docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py:135–174

    @I.ir_module
    class DebugModule(BasePyModule):
        @T.prim_func(s_tir=True)
        def matmul_tir(var_A: T.handle, var_B: T.handle, var_C: T.handle):
            n = T.int32()
            A = T.match_buffer(var_A, (n, 4), "float32")
            B = T.match_buffer(var_B, (4, 3), "float32")
            C = T.match_buffer(var_C, (n, 3), "float32")
            for i, j, k in T.grid(n, 3, 4):
                with T.sblock("matmul"):
                    vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                    with T.init():
                        C[vi, vj] = T.float32(0)
                    C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]

        @I.pyfunc
        def forward(self, x, weights):
            # Inspect input
            print(f"  [DEBUG] input shape: {x.shape}, mean: {x.mean():.4f}")

            # Run TIR matmul
            x_tvm = self._convert_pytorch_to_tvm(x)
            w_tvm = self._convert_pytorch_to_tvm(weights)
            out = self.call_tir(
                self.matmul_tir,
                [x_tvm, w_tvm],
                out_sinfo=R.Tensor((x.shape[0], 3), "float32"),
            )
            logits = self._convert_tvm_to_pytorch(out)

            # Inspect intermediate value — impossible with a compiled-only workflow
            print(
                f"  [DEBUG] logits shape: {logits.shape}, "
                f"min: {logits.min():.4f}, max: {logits.max():.4f}"
            )

            result = F.softmax(logits, dim=-1)

            # Verify output
            print(f"  [DEBUG] probs sum: {result.sum(dim=-1)}")
            return result

    mod = DebugModule(device=tvm.cpu(0))