debug

import torch

import tvm
from tvm import tir
from tvm.relax.frontend import nn
from tvm.relax.frontend.nn import op, spec
from tvm.runtime import NDArray

print

class Layer(nn.Module):
    def forward(self, x: nn.Tensor):  # pylint: disable=invalid-name
        op.print_(x)
        return x

model = Layer().jit(
    spec={
        "forward": {"x": spec.Tensor([10, 5], dtype="float32")},
    },
    debug=True,
)
x = torch.rand((10, 5), dtype=torch.float32)  # pylint: disable=invalid-name
y = model["forward"](x)  # pylint: disable=invalid-name
assert isinstance(y, torch.Tensor)

/tmp/ipykernel_1225526/3381093976.py:3: shape = (10, 5), dtype = float32, data =
[[0.73462075 0.01047635 0.31426018 0.8674046  0.12726694]
 [0.05869561 0.14252067 0.57591957 0.62454313 0.5909737 ]
 [0.04823518 0.8938983  0.89632434 0.15392512 0.18556875]
 [0.31064558 0.6404101  0.62769175 0.1271398  0.89232284]
 [0.16888297 0.4059453  0.70379585 0.77904415 0.84397256]
 [0.00197196 0.48467797 0.3471346  0.00262749 0.82645226]
 [0.28290933 0.51377964 0.67835134 0.04419875 0.07198626]
 [0.00947952 0.96474004 0.16065902 0.33178765 0.49370688]
 [0.6731469  0.5020207  0.44368893 0.4925446  0.45231378]
 [0.48918706 0.0983749  0.9165322  0.5107241  0.788272  ]]

调试函数

@tvm.register_func("testing.relax.frontend.nn.test_debug_func")
def _debug(  # pylint: disable=too-many-arguments
    lineno: str,
    tensor: NDArray,
    const_int: int,
    const_float: float,
    const_str: str,
    var_int: int,
) -> None:
    assert "test_frontend_nn_debug.py" in lineno
    assert tensor.shape == (10, 5)
    assert const_int == 1
    assert const_float == 2.0
    assert const_str == "test"
    assert var_int == 8

class Layer(nn.Module):
    def forward(self, x: nn.Tensor, v: tir.Var):  # pylint: disable=invalid-name
        op.debug_func("testing.relax.frontend.nn.test_debug_func", x, 1, 2.0, "test", v)
        return x

model = Layer().jit(
    spec={
        "forward": {
            "x": spec.Tensor([10, 5], dtype="float32"),
            "v": "int",
        },
    },
    debug=True,
)
x = torch.rand((10, 5), dtype=torch.float32)  # pylint: disable=invalid-name
y = model["forward"](x, 8)  # pylint: disable=invalid-name
assert isinstance(y, torch.Tensor)

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[4], line 32
     22 model = Layer().jit(
     23     spec={
     24         "forward": {
   (...)
     29     debug=True,
     30 )
     31 x = torch.rand((10, 5), dtype=torch.float32)  # pylint: disable=invalid-name
---> 32 y = model["forward"](x, 8)  # pylint: disable=invalid-name
     33 assert isinstance(y, torch.Tensor)

File /media/pc/data/lxw/ai/tvm/python/tvm/relax/frontend/nn/torch.py:78, in TorchModule.__getitem__.<locals>._closure(*args)
     71 args = [
     72     _torch_to_tvm(arg_name, arg_spec, arg)
     73     for arg_name, arg_spec, arg in zip(
     74         method_spec.arg_names, method_spec.arg_specs, args
     75     )
     76 ]
     77 if self.effects is not None:
---> 78     outputs, self.effects = method(*args, *self.effects, *self.params)
     79 else:
     80     outputs = method(*args, *self.params)

File /media/pc/data/lxw/ai/tvm/python/tvm/_ffi/_cython/packed_func.pxi:339, in tvm._ffi._cy3.core.PackedFuncBase.__call__()

File /media/pc/data/lxw/ai/tvm/python/tvm/_ffi/_cython/packed_func.pxi:270, in tvm._ffi._cy3.core.FuncCall()

File /media/pc/data/lxw/ai/tvm/python/tvm/_ffi/_cython/packed_func.pxi:259, in tvm._ffi._cy3.core.FuncCall3()

File /media/pc/data/lxw/ai/tvm/python/tvm/_ffi/_cython/base.pxi:185, in tvm._ffi._cy3.core.CHECK_CALL()

File /media/pc/data/lxw/ai/tvm/python/tvm/_ffi/base.py:468, in raise_last_ffi_error()
    462 # The exception PyObject may contain a large amount of state,
    463 # including all stack frames that may be inspected in a later
    464 # PDB post-mortem.  Therefore, we must make sure to remove the
    465 # underlying PyObject* from the C++ side after we retrieve it.
    466 _LIB.TVMDropLastPythonError()
--> 468 raise py_err

File /media/pc/data/lxw/ai/tvm/src/runtime/relax_vm/vm.cc:558, in tvm::runtime::relax_vm::VirtualMachineImpl::InvokeClosurePacked(tvm::runtime::ObjectRef const&, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)()
    556   {
    557     NVTXScopedRange scope("RelaxVM: " + clo->func_name);
--> 558     clo->impl.CallPacked(TVMArgs(values.data(), tcodes.data(), args.size() + 1), rv);
    559   }
    560 }

File /media/pc/data/lxw/ai/tvm/src/runtime/relax_vm/vm.cc:632, in operator()()
    630     inputs[i] = args[i + 1];
    631   }
--> 632   *rv = static_cast<VirtualMachineImpl*>(ctx_ptr)->InvokeBytecode(gf_idx, inputs);
    633 });
    634 return VMClosure(func_name, impl);

File /media/pc/data/lxw/ai/tvm/src/runtime/relax_vm/vm.cc:703, in tvm::runtime::relax_vm::VirtualMachineImpl::InvokeBytecode(long, std::vector<tvm::runtime::TVMRetValue, std::allocator<tvm::runtime::TVMRetValue> > const&)()
    701   // set program counter
    702   pc_ = gfunc.start_instr;
--> 703   RunLoop();
    704   return return_value_;
    705 }

File /media/pc/data/lxw/ai/tvm/src/runtime/relax_vm/vm.cc:828, in tvm::runtime::relax_vm::VirtualMachineImpl::RunLoop()()
    826 switch (instr.op) {
    827   case Opcode::Call: {
--> 828     this->RunInstrCall(curr_frame, instr);
    829     break;
    830   }

File /media/pc/data/lxw/ai/tvm/src/runtime/relax_vm/vm.cc:781, in tvm::runtime::relax_vm::VirtualMachineImpl::RunInstrCall(tvm::runtime::relax_vm::VMFrame*, tvm::runtime::relax_vm::Instruction)()
    779 
    780   if (instrument_ == nullptr) {
--> 781     this->InvokeClosurePacked(func_pool_[instr.func_idx], args, &ret);
    782   } else {
    783     // insert light-weight instrument callback

File /media/pc/data/lxw/ai/tvm/src/runtime/relax_vm/builtin.cc:497, in operator()()
    495     }
    496   }
--> 497   debug_func->CallPacked(TVMArgs(call_args.data(), call_type_codes.data(), num_args + 1), rv);
    498   *rv = io_effect;
    499 });

File /media/pc/data/lxw/ai/tvm/python/tvm/_ffi/_cython/packed_func.pxi:56, in tvm._ffi._cy3.core.tvm_callback()

Cell In[4], line 10, in _debug(lineno, tensor, const_int, const_float, const_str, var_int)
      1 @tvm.register_func("testing.relax.frontend.nn.test_debug_func")
      2 def _debug(  # pylint: disable=too-many-arguments
      3     lineno: str,
   (...)
      8     var_int: int,
      9 ) -> None:
---> 10     assert "test_frontend_nn_debug.py" in lineno
     11     assert tensor.shape == (10, 5)
     12     assert const_int == 1

AssertionError: