Manually remove operation from graph

Development
#1

I’m working on BYOC, and suspect that a particular operation is responsible for an error.

Given a mod (i.e., of type tvm.ir.module.IRModule), is there a way to forceibly remove a given operation, so that it is not compiled with the rest of the model? In my case it is the last operation, so we do not need to worry about dependencies.

E.g., if I have a 4 op mod, can I drop the statement which generates %4 before compilation?

def @main(%serving_default_input_1:0: Tensor[(1, 32, 32, 16), int8] /* ty=Tensor[(1, 32, 32, 16), int8] span=serving_default_input_1:0:0:0 */, %v_param_1: Tensor[(3, 3, 16, 16), int8] /* ty=Tensor[(3, 3, 16, 16), int8] span=model/conv2d/Co
nv2D:0:0 */, %v_param_2: Tensor[(16), int32] /* ty=Tensor[(16), int32] */, %v_param_3: Tensor[(3, 3, 16, 1), int8] /* ty=Tensor[(3, 3, 16, 1), int8]  */,
 %v_param_4: Tensor[(16), int32] /* ty=Tensor[(16), int32]  */, output_tensor_names=["StatefulPartitionedCall_0"]) -> Tensor[(1, 28, 28, 16), int8] {
  %0 = qnn.conv2d(%serving_default_input_1:0, %v_param_1, 6 /* ty=int32 */, 0 /* ty=int32 */, 0.0409262f /* ty=float32  */, meta[relay.Constant][0] /* ty=Tensor[(16), flo
at32] */, padding=[0, 0, 0, 0], channels=16, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO", out_dtype="int32") /* ty=Tens
or[(1, 30, 30, 16), int32]  */;
  %1 = add(%0, %v_param_2) /* ty=Tensor[(1, 30, 30, 16), int32] */;
  %2 = qnn.requantize(%1, meta[relay.Constant][1] /* ty=Tensor[(16), float32]  */, 0 /* ty=int32 */, 0.0214691f /* ty=float32 */, -128 /* ty=int32 */, axis=3, out_dtype="int8") /* ty=Tensor[(1, 30, 30, 16), int8] */;
  %3 = @tvmgen_default_n64_main_0(%2, %v_param_3) /* ty=Tensor[(1, 28, 28, 16), int32] */;
  %4 = add(%3, %v_param_4) /* ty=Tensor[(1, 28, 28, 16), int32] */;
  qnn.requantize(%4, meta[relay.Constant][2] /* ty=Tensor[(16), float32]  */, 0 /* ty=int32 */, 0.0117953f /* ty=float32 */, -9 /* ty=int32 span=StatefulPartitionedCall:0:0:0 */, axis=3, out_dtype="int8") /* ty=Tensor[(1, 28, 28, 16), int8] */
}
#2

Hi,

It is definitely possible to remove the last operation. The following quick and dirty solution worked for me (if you’re certain the last node in the body of main is a Call and you want to keep its first argument)

def remove_last_op(mod: tvm.IRModule):
    main = mod['main']
    new_body = main.body.args[0]
    return tvm.IRModule.from_expr(new_body)

Which turns

def @main(%x: Tensor[(1, 32, 4, 4), float32]) {
  %0 = nn.conv2d(%x, meta[relay.Constant][0], padding=[0, 0, 0, 0]);
  nn.bias_add(%0, meta[relay.Constant][1])
}

into

def @main(%x: Tensor[(1, 32, 4, 4), float32]) {
  nn.conv2d(%x, meta[relay.Constant][0], padding=[0, 0, 0, 0])
}
1 Like
#3

Thanks, I’ll add this to my bag of tricks. I’ve not modified IRModule bodies in this way before, good to know this is how it’s done.