Different Results When Applying The Same Relay Transform: FoldScaleAxis

jeff815 · October 29, 2021, 4:27pm

I am observing different results using the tvm.relay.transform.FoldScaleAxis pass depending on the method I use to call the pass.

Calling the pass in the manner:

new_mod = relay.transform.FoldScaleAxis()(mod)

Differs from invoking it in the following manners (both of which produce equivalent results):

with tvm.transform.PassContext(opt_level=3):
    new_mod = relay.transform.FoldScaleAxis()(mod)

seq = tvm.ir.transform.Sequential([relay.transform.FoldScaleAxis()])
with tvm.transform.PassContext(opt_level=3):
    new_mod = seq(mod)

Is this expected behavior?

The following is example code demonstrating the behavior:

import tvm
from tvm import relay
import numpy as np

def create_model():
    x = relay.var('x', shape=(1, 32, 10, 10), dtype='float32')
    pool = relay.nn.adaptive_avg_pool2d(x, output_size=(1, 1))
    relu = relay.nn.relu(pool)
    mul = relay.multiply(x, relu)
    conv = relay.nn.conv2d(mul, relay.const(np.random.rand(16, 32, 1, 1).astype('float32')), kernel_size=(1, 1), channels=16)
    mod = tvm.IRModule.from_expr(conv)
    return mod

mod = create_model()

mod_0 = relay.transform.FoldScaleAxis()(mod)

with tvm.transform.PassContext(opt_level=3):
    mod_1 = relay.transform.FoldScaleAxis()(mod)

seq = tvm.ir.transform.Sequential([relay.transform.FoldScaleAxis()])
with tvm.transform.PassContext(opt_level=3):
    mod_2 = seq(mod)

assert not tvm.ir.structural_equal(mod_0, mod_1)

assert tvm.ir.structural_equal(mod_1, mod_2)

JoeyChou · October 29, 2021, 5:06pm

Kindly tag @comaniac. Can you check this post if you have some time please?

comaniac · October 29, 2021, 5:15pm

The first method doesn’t specify any pass context so FoldScaleAxis won’t be applied, because two passes in its sequence are opt_level=3 passes:

github.com

apache/tvm/blob/main/src/relay/transforms/fold_scale_axis.cc#L1012




}  // namespace fold_scale_axis


namespace transform {


Pass ForwardFoldScaleAxis() {
  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
      [=](Function f, IRModule m, PassContext pc) {
        return Downcast<Function>(relay::fold_scale_axis::ForwardFoldScaleAxis(f));
      };
  return CreateFunctionPass(pass_func, 3, "ForwardFoldScaleAxis", {"InferType"});
}


TVM_REGISTER_GLOBAL("relay._transform.ForwardFoldScaleAxis").set_body_typed(ForwardFoldScaleAxis);


Pass BackwardFoldScaleAxis() {
  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
      [=](Function f, IRModule m, PassContext pc) {
        return Downcast<Function>(relay::fold_scale_axis::BackwardFoldScaleAxis(f));
      };
  return CreateFunctionPass(pass_func, 3, "BackwardFoldScaleAxis", {"InferType"});

github.com

apache/tvm/blob/main/src/relay/transforms/fold_scale_axis.cc#L1022


  return CreateFunctionPass(pass_func, 3, "ForwardFoldScaleAxis", {"InferType"});
}


TVM_REGISTER_GLOBAL("relay._transform.ForwardFoldScaleAxis").set_body_typed(ForwardFoldScaleAxis);


Pass BackwardFoldScaleAxis() {
  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
      [=](Function f, IRModule m, PassContext pc) {
        return Downcast<Function>(relay::fold_scale_axis::BackwardFoldScaleAxis(f));
      };
  return CreateFunctionPass(pass_func, 3, "BackwardFoldScaleAxis", {"InferType"});
}


TVM_REGISTER_GLOBAL("relay._transform.BackwardFoldScaleAxis").set_body_typed(BackwardFoldScaleAxis);


Pass FoldScaleAxis() {
  // FoldScaleAxis pass contains the following three passes. Therefore, we can
  // register it as a sequential pass.
  Pass pass = Sequential({BackwardFoldScaleAxis(), ForwardFoldScaleAxis(), FoldConstant()},
                         "FoldScaleAxis");
  return pass;

Without opt_level=3 in the pass context, only the FoldConstant will be applied, because it is an opt_level=2 pass:

github.com

apache/tvm/blob/main/src/relay/transforms/fold_constant.cc#L366



TVM_REGISTER_GLOBAL("relay._transform.FoldConstantExpr").set_body_typed(FoldConstant);

namespace transform {

Pass FoldConstant() {
  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
      [=](Function f, IRModule m, PassContext pc) {
        return Downcast<Function>(FoldConstant(f, m));
      };
  return CreateFunctionPass(pass_func, 2, "FoldConstant", {});
}

TVM_REGISTER_GLOBAL("relay._transform.FoldConstant").set_body_typed(FoldConstant);

}  // namespace transform

}  // namespace relay
}  // namespace tvm

jeff815 · October 29, 2021, 5:27pm

@comaniac

Thanks for the explanation. In terms of expected behavior, the result is different when I use other passes like FoldConstant (which is an opt_level=2 pass). Regardless of which method is used the FoldConstant pass applies the same transformation. Should FoldScaleAxis be altered so it behaves in the same way?

Here is an example for FoldConstant:

import tvm
from tvm import relay
import numpy as np

def create_model():
    data = tvm.nd.array(np.array([1.0, 2.0, 3.0]))
    const = relay.const(data)
    concat = relay.op.concatenate([const, const], axis=0)
    mod = tvm.IRModule.from_expr(concat)
    return mod

mod = create_model()

mod_0 = relay.transform.FoldConstant()(mod)

with tvm.transform.PassContext(opt_level=2):
    mod_1 = relay.transform.FoldConstant()(mod)

seq = tvm.ir.transform.Sequential([relay.transform.FoldConstant()])
with tvm.transform.PassContext(opt_level=2):
    mod_2 = seq(mod)

assert tvm.ir.structural_equal(mod_0, mod_1)

assert tvm.ir.structural_equal(mod_1, mod_2)

comaniac · October 29, 2021, 5:38pm

I don’t think so. FoldConstant behaves the same for both methods just because it is opt_level=2, which is the default level. Passes such as FoldScaleAxis are assigned to higher level for reasons, so you should always use the PassContext to specify the desire opt_level when invoking passes.

jeff815 · October 29, 2021, 6:51pm

In the case of CanonicalizeOps which has an opt_level=3 we get the same behavior as FoldConstant where we don’t have to specify the opt level for it to work in the case of mod_0 = relay.transform.CanonicalizeOps()(mod)

Example Code:

import tvm
from tvm import relay
import numpy as np

def create_model():
    x = relay.var("x", shape=(1, 64, 56, 56))
    bias = relay.var("bias", shape=(64,))
    weight = relay.var("weight", shape=(64, 64, 3, 3))
    conv = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1))
    bias_add = relay.nn.bias_add(conv, bias)
    mod = tvm.IRModule.from_expr(bias_add)
    return mod

mod = create_model()
mod = relay.transform.InferType()(mod)

mod_0 = relay.transform.CanonicalizeOps()(mod)

with tvm.transform.PassContext(opt_level=3):
    mod_1 = relay.transform.CanonicalizeOps()(mod)

seq = tvm.ir.transform.Sequential([relay.transform.CanonicalizeOps()])
with tvm.transform.PassContext(opt_level=3):
    mod_2 = seq(mod)

assert tvm.ir.structural_equal(mod_0, mod_1)

assert tvm.ir.structural_equal(mod_1, mod_2)

comaniac · October 29, 2021, 7:48pm

Sorry my previous description wasn’t accurate. If you directly invoke a pass such as relay.transform.FoldConstant()(mod), then it may or may not consider the current opt_level. It depends on whether the pass is a function pass or a sequence pass.

For example, FoldConstant and CanonicalizeOps are function passes. In this case, if you invoke them directly, it will be applied for sure.

github.com

apache/tvm/blob/main/src/relay/transforms/canonicalize_ops.cc#L78


  return PostOrderRewrite(e, &rewriter);
}


namespace transform {


Pass CanonicalizeOps() {
  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
      [=](Function f, IRModule m, PassContext pc) {
        return Downcast<Function>(CanonicalizeOps(f));
      };
  return CreateFunctionPass(pass_func, 3, "CanonicalizeOps", {"InferType"});
}


TVM_REGISTER_GLOBAL("relay._transform.CanonicalizeOps").set_body_typed(CanonicalizeOps);


}  // namespace transform


}  // namespace relay
}  // namespace tvm

On the other hand, FoldScaleAxis is a sequence pass:

github.com

apache/tvm/blob/main/src/relay/transforms/fold_scale_axis.cc#L1030


        return Downcast<Function>(relay::fold_scale_axis::BackwardFoldScaleAxis(f));
      };
  return CreateFunctionPass(pass_func, 3, "BackwardFoldScaleAxis", {"InferType"});
}


TVM_REGISTER_GLOBAL("relay._transform.BackwardFoldScaleAxis").set_body_typed(BackwardFoldScaleAxis);


Pass FoldScaleAxis() {
  // FoldScaleAxis pass contains the following three passes. Therefore, we can
  // register it as a sequential pass.
  Pass pass = Sequential({BackwardFoldScaleAxis(), ForwardFoldScaleAxis(), FoldConstant()},
                         "FoldScaleAxis");
  return pass;
}


TVM_REGISTER_GLOBAL("relay._transform.FoldScaleAxis").set_body_typed(FoldScaleAxis);


}  // namespace transform


}  // namespace relay
}  // namespace tvm

It means the following two methods are equivalent, and both methods are affected by the current opt_level.

mod_0 = relay.transform.FoldScaleAxis()(mod)

seq = tvm.ir.transform.Sequential([relay.transform.FoldScaleAxis()])
mod_2 = seq(mod)

Again, it is recommended to always use PassContext with a pass sequence to run the required passes.

jeff815 · November 8, 2021, 6:03pm

Do you think it would be a good idea to modify the description of the FoldScaleAxis pass? I would add that that FoldScaleAxis is a sequential pass and must be placed within a PassContext of 3 for all passes contained within it to execute.

comaniac · November 8, 2021, 6:26pm

I would instead prefer one of the following strategies:

We don’t expect users to call passes individually (no change).
Think of an approach to enforce the pass execution when running it without an outside pass sequential.

cc @zhiics @jroesch @tqchen for comments.