Question concerning assignment 4, with respect to `call_tir` #31

Ja1Zhou · 2022-07-25T11:16:31Z

Ja1Zhou
Jul 25, 2022

Hi, I was trying to work out assignment 4. In section 4, we are instructed to register an external function for conv2d and incorporate it into BlockBuilder using BlockBuilder.emit. In the markdown file, it was hinted that we should also use call_tir. However, I find this part of description quite vague and I am not quite sure how to achieve this. The best I could think of is lv0 = bb.emit(relax.call_tir("env.conv2d", (x, conv2d_weight, conv2d_bias), (1, 32, 26, 26), dtype="float32")). I indeed inspected the source code and documents, yet I couldn't find the answers. Could you please elaborate on this point further? Many thanks!

Answered by jinhongyii

Jul 25, 2022

You need to use relax.extern("env.conv2d") to get the extern function (its type is Expr). I have to admit this pattern of emitting extern function is complex and we are trying to figure out a way to make it simpler and easier to use.

View full answer

jinhongyii · 2022-07-25T11:39:49Z

jinhongyii
Jul 25, 2022
Maintainer

Sorry for the inconvenience. You are really close to the answer. To insert a call_tir using BlockBuilder, you should use relax.op.call_tir instead of relax.call_tir. relax.op.call_tir is a function defined in https://github.com/mlc-ai/relax/blob/relax/python/tvm/relax/op/base.py#L27 which returns a call node for the call_tir operator.

1 reply

Ja1Zhou Jul 25, 2022
Author

Hi, thanks for helping out! I was somehow still unable to get past this error.

  File "assignment.py", line 260, in <module>
    mod = create_model_with_torch_func()
  File "assignment.py", line 245, in create_model_with_torch_func
    lv0 = bb.emit(relax.op.call_tir("env.conv2d", (x, conv2d_weight, conv2d_bias), (1, 32, 26, 26), dtype="float32"))
  File "/home/zhejian/miniconda3/envs/mlc/lib/python3.8/site-packages/tvm/relax/op/base.py", line 74, in call_tir
    return _ffi_api.call_tir(func, args, shape, output_type, tir_vars)
  File "tvm/_ffi/_cython/./packed_func.pxi", line 331, in tvm._ffi._cy3.core.PackedFuncBase.__call__
  File "tvm/_ffi/_cython/./packed_func.pxi", line 276, in tvm._ffi._cy3.core.FuncCall
  File "tvm/_ffi/_cython/./base.pxi", line 181, in tvm._ffi._cy3.core.CHECK_CALL
tvm._ffi.base.TVMError: Traceback (most recent call last):
  4: TVMFuncCall
  3: _ZN3tvm7runtime13PackedFun
  2: tvm::runtime::TypedPackedFunc<tvm::RelayExpr (tvm::RelayExpr, tvm::relay::Tuple, tvm::RelayExpr, tvm::Type, tvm::runtime::Optional<tvm::RelayExpr>)>::AssignTypedLambda<tvm::RelayExpr (*)(tvm::RelayExpr, tvm::relay::Tuple, tvm::RelayExpr, tvm::Type, tvm::runtime::Optional<tvm::RelayExpr>)>(tvm::RelayExpr (*)(tvm::RelayExpr, tvm::relay::Tuple, tvm::RelayExpr, tvm::Type, tvm::runtime::Optional<tvm::RelayExpr>), std::string)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const
  1: tvm::runtime::TVMMovableArgValueWithContext_::operator tvm::RelayExpr<tvm::RelayExpr>() const
  0: _ZN3tvm7runtime6deta
  5: TVMFuncCall
  4: _ZN3tvm7runtime13PackedFun
  3: tvm::runtime::TypedPackedFunc<tvm::RelayExpr (tvm::RelayExpr, tvm::relay::Tuple, tvm::RelayExpr, tvm::Type, tvm::runtime::Optional<tvm::RelayExpr>)>::AssignTypedLambda<tvm::RelayExpr (*)(tvm::RelayExpr, tvm::relay::Tuple, tvm::RelayExpr, tvm::Type, tvm::runtime::Optional<tvm::RelayExpr>)>(tvm::RelayExpr (*)(tvm::RelayExpr, tvm::relay::Tuple, tvm::RelayExpr, tvm::Type, tvm::runtime::Optional<tvm::RelayExpr>), std::string)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const
  2: tvm::runtime::TVMMovableArgValueWithContext_::operator tvm::RelayExpr<tvm::RelayExpr>() const
  1: tvm::RelayExpr tvm::runtime::TVMPODValue_::AsObjectRef<tvm::RelayExpr>() const
  0: _ZN3tvm7runtime6deta
  File "/workspace/tvm/include/tvm/runtime/packed_func.h", line 777
TVMError: In function relax.op.call_tir(0: RelayExpr, 1: relay.Tuple, 2: RelayExpr, 3: Type, 4: RelayExpr) -> RelayExpr: error while converting argument 0: [19:50:44] /workspace/tvm/include/tvm/runtime/packed_func.h:1887: 
---------------------------------------------------------------
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
---------------------------------------------------------------

  Check failed: type_code_ == kTVMObjectHandle (11 vs. 8) : expected Object but got str

It seems that relax.op.call_tir accepts an Expr instead of str. I believe I am not aware of the differences between relax.op.call_tir and tvm.script.relax.

My relevant code is as follows:

# %%
# register pytorch conv2d function
@tvm.register_func("env.conv2d", override=True)
def torch_conv2d(
    x: tvm.nd.NDArray,
    w: tvm.nd.NDArray,
    b: tvm.nd.NDArray,
    out: tvm.nd.NDArray
):
    x_torch = torch.from_dlpack(x)
    w_torch = torch.from_dlpack(w)
    b_torch = torch.from_dlpack(b)
    out_torch = torch.from_dlpack(out)
    out_torch = torch.nn.functional.conv2d(x_torch, w_torch, b_torch)
# %%
# function for mixed pytorch + topi model
def mixed_pytorch_model(A, w4, b4, w6, b6):
    B = topi.nn.relu(A)
    B = topi.nn.pool2d(B, [2,2], [2,2], [1,1], padding=[0,0,0,0], pool_type="max")
    B = topi.nn.flatten(B)
    B = topi.matmul(B, w4, transp_b=True)

    B = topi.add(B, b4)
    B = topi.nn.relu(B)
    B = topi.matmul(B, w6, transp_b=True)
    
    B = topi.add(B, b6)
    B = topi.nn.softmax(B)
    return B
# %%
# function to create IR with mixed model
def create_model_with_torch_func():
    bb = relax.BlockBuilder()

    x = relax.Var("x", input_shape, relax.DynTensorType(4, "float32"))

    conv2d_weight = relax.const(weight_map["conv2d_weight"], "float32")
    conv2d_bias = relax.const(weight_map["conv2d_bias"].reshape(1, 32, 1, 1), "float32")
    linear0_weight = relax.const(weight_map["linear0_weight"], "float32")
    linear0_bias = relax.const(weight_map["linear0_bias"].reshape(1, 100), "float32")
    linear1_weight = relax.const(weight_map["linear1_weight"], "float32")
    linear1_bias = relax.const(weight_map["linear1_bias"].reshape(1, 10), "float32")
    
    with bb.function("main", [x]):
        with bb.dataflow():
            lv0 = bb.emit(relax.op.call_tir("env.conv2d", (x, conv2d_weight, conv2d_bias), (1, 32, 26, 26), dtype="float32"))
            lv1 = bb.emit_te(
                mixed_pytorch_model, x,
                linear0_weight,
                linear0_bias,
                linear1_weight,
                linear1_bias
            )
            gv = bb.emit_output(lv1)
        bb.emit_func_output(gv)

    return bb.get()


test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=False)
mod = create_model_with_torch_func()
check_equivalence(mod, torch_model, test_loader)

Thanks again for helping out!

jinhongyii · 2022-07-25T12:10:06Z

jinhongyii
Jul 25, 2022
Maintainer

You need to use relax.extern("env.conv2d") to get the extern function (its type is Expr). I have to admit this pattern of emitting extern function is complex and we are trying to figure out a way to make it simpler and easier to use.

1 reply

Ja1Zhou Jul 25, 2022
Author

Problem solved, thanks a lot~

GeLee-Q · 2022-07-27T07:56:40Z

GeLee-Q
Jul 27, 2022

Hi，In the Section 3 of assignment 4, the model use torch.conv2d and tvm.nn has the following error. Have you encountered this problem?

my_code

@tvm.register_func("env.conv2d", override=True)
def torch_conv2d(x: tvm.nd.NDArray, 
          w: tvm.nd.NDArray, 
          b: tvm.nd.NDArray,
          out: tvm.nd.NDArray ):
  x_torch = torch.from_dlpack(x)
  w_torch = torch.from_dlpack(w)
  b_torch = torch.from_dlpack(b)
  out_torch = torch.from_dlpack(out)

  out_torch = torch.nn.functional.conv2d(x_torch, w_torch)
  out_torch = torch.add(out_torch, b_torch)



def conv2d_1(Input, fliter, bias):
  lv1_0 = tvm.topi.nn.conv2d(Input, fliter, 1, 0, 1)
  return tvm.topi.add(lv1_0, bias)
  
def relu_2(Input):
  tvm.topi.nn.pool2d
  return tvm.topi.nn.relu(Input)

def maxPool_3(Input):
  # return tvm.topi.nn.pool2d(Input, (2, 2), (0, 0) , 1, (0, 0, 0, 0), 'max', False, 'NHCW', False)
  return tvm.topi.nn.pool2d(data = Input, kernel = [2, 2], dilation = (1,1), stride = [2,2], padding = [0,0,0,0], pool_type = 'max')
   
def flatten_4(Input):
  return tvm.topi.nn.flatten(Input)

def linear_5(Input, weight, bias):
  lv5_0 = tvm.topi.nn.dense(Input, weight)
  return tvm.topi.add(lv5_0, bias) 
    
def relu_6(Input):
  return tvm.topi.nn.relu(Input)

def linear_7(Input, weight ,bias):
  lv7_0 = tvm.topi.nn.dense(Input, weight)
  return tvm.topi.add(lv7_0, bias) 

def softMax_8(Intput):
  return tvm.topi.nn.softmax(Intput, axis=- 1)

# 返回构造得到的IRModule

def create_model_with_torch_func():
    bb = relax.BlockBuilder()
    x = relax.Var("x", input_shape, relax.DynTensorType(batch_size, "float32"))
    conv2d_weight = relax.const(weight_map["conv2d_weight"], "float32")
    conv2d_bias = relax.const(weight_map["conv2d_bias"].reshape(1, 32, 1, 1), "float32")
    linear0_weight = relax.const(weight_map["linear0_weight"], "float32")
    linear0_bias = relax.const(weight_map["linear0_bias"].reshape(1, 100), "float32")
    linear1_weight = relax.const(weight_map["linear1_weight"], "float32")
    linear1_bias = relax.const(weight_map["linear1_bias"].reshape(1, 10), "float32")

    with bb.function("main", [x]):
        with bb.dataflow():
          lv1 = bb.emit(relax.op.call_tir(relax.extern("env.conv2d"), (x, conv2d_weight, conv2d_bias),(batch_size, 32, 26, 26), dtype="float32"))
          lv2 = bb.emit_te(relu_2, lv1)
          lv3 = bb.emit_te(maxPool_3, lv2)
          lv4 = bb.emit_te(flatten_4, lv3)
          lv5 = bb.emit_te(linear_5, lv4, linear0_weight, linear0_bias)
          lv6 = bb.emit_te(relu_6, lv5)
          lv7 = bb.emit_te(linear_7, lv6, linear1_weight, linear1_bias)
          lv8 = bb.emit_te(softMax_8, lv7)

          gv = bb.emit_output(lv8)
        bb.emit_func_output(gv)

    return bb.get()


def check_equivalence(mod, torch_model, test_loader):
    torch_model.eval()
    with torch.no_grad():
        rt_mod = build_mod(mod)
        for data, label in test_loader:
            data, label = data.cpu(), label.cpu()
            output_from_pytorch = torch_model(data).numpy()
            output_from_relax = rt_mod["main"](tvm.nd.array(data, tvm.cpu())).numpy()
            # print(output_from_pytorch)
            # print('---------------------------')
            # print(output_from_relax)
            tvm.testing.assert_allclose(output_from_pytorch, output_from_relax, rtol=1e-4, atol=1e-7)


mod = create_model_with_torch_func()
# IPython.display.HTML(code2html(mod.script()))

test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=False)
check_equivalence(mod, torch_model, test_loader)

6 replies

GeLee-Q Jul 27, 2022

I'm sorry that it's hard for me to discover the answer behind the hint.
If troch.nn.Conv2d and torch.nn.functional.conv2d are just the difference in calling form, why does this problem occur?

@tvm.register_func("env.conv2d", override=True)
def torch_conv2d(
    x: tvm.nd.NDArray,
    w: tvm.nd.NDArray,
    out : tvm.nd.NDArray
):
    x_torch = torch.from_dlpack(x)
    w_torch = torch.from_dlpack(w)
    out_torch = torch.from_dlpack(out)
    out_torch = torch.nn.functional.conv2d(x_torch, w_torch, padding=0, stride = 1, dilation = 1)
    out = out_torch

def add(Input, bias):
  return tvm.topi.add(Input, bias) 


// bb.dataflow
lv1_0 = bb.emit(relax.op.call_tir(relax.extern("env.conv2d"), (x, conv2d_weight),(4, 32, 26, 26), dtype="float32"))
lv1 = bb.emit_te(add, lv1_0, conv2d_bias)

Ja1Zhou Jul 27, 2022
Author

I worked around this problem by specifying an add function with 0 after convolution, and specifying out=out_torch as in the given example. I believe calling a copy construct of tensor is changing the underling tensor memory.

GeLee-Q Jul 28, 2022

I have some confusion about "specifying an add function with 0 after convolution"
The following code doesn't work

out_torch = torch.nn.functional.conv2d(x_torch, w_torch)
zero = torch.Tensor(86528*[0]).reshape(-1, 32, 26, 26)
torch.add(out_torch, zero, out = out_torch)

yuko29 Jul 28, 2022

I have the same problem with @GeLee-Q . I found the point is that torch.nn.functional.conv2d will return the new output tensor.
If you write out_torch = torch.nn.functional.conv2d(...), the out_torch at left hand side will be that new output tensor. It means that the variable out_torch will no longer point to same the memory space as out as expected.

My solution is to copy the output tensor that torch.nn.functional.conv2d returns to the space that out points to.

t = torch.nn.functional.conv2d(x_torch, w_torch)
out_torch.copy_(t)

It can pass the equivalence check, but it costs more memory spaces.

GeLee-Q Jul 29, 2022

Thanks! The memory copy construct is like cpp's memory behavior....

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Question concerning assignment 4, with respect to `call_tir` #31

{{title}}

Replies: 3 comments 8 replies

{{title}}

{{title}}

{{title}}

{{editor}}'s edit

{{editor}}'s edit

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

Select a reply

Question concerning assignment 4, with respect to call_tir #31

Ja1Zhou Jul 25, 2022

Replies: 3 comments · 8 replies

jinhongyii Jul 25, 2022 Maintainer

Ja1Zhou Jul 25, 2022 Author

jinhongyii Jul 25, 2022 Maintainer

Ja1Zhou Jul 25, 2022 Author

GeLee-Q Jul 27, 2022

GeLee-Q Jul 27, 2022

Ja1Zhou Jul 27, 2022 Author

GeLee-Q Jul 28, 2022

yuko29 Jul 28, 2022

GeLee-Q Jul 29, 2022

Question concerning assignment 4, with respect to `call_tir` #31

Ja1Zhou
Jul 25, 2022

Replies: 3 comments 8 replies

jinhongyii
Jul 25, 2022
Maintainer

Ja1Zhou Jul 25, 2022
Author

jinhongyii
Jul 25, 2022
Maintainer

Ja1Zhou Jul 25, 2022
Author

GeLee-Q
Jul 27, 2022

Ja1Zhou Jul 27, 2022
Author