[PatternLang] Add ConstantPattern

[comaniac]

Discussion: https://discuss.tvm.ai/t/patternlang-match-constant-nodes/6835

Changelog:

• Add ConstantPattern to the pattern language.
• Add unit tests of constant pattern.
• Eliminate pylint errors in the dataflow pattern unit test file.
• Update the document.

cc @mbrookhart @zhiics @mbaret

[mbrookhart]

To make other syntactic sugar, you could add an is_const() that simiply creats a new ConstPattern, but I don't think it's strictly necessary.

For my own edification, what did you use to auto-format the test file? I tried autopep8 on it a while back and got really bad results, so I ignored it after that (as does the linter).

[comaniac]

To make other syntactic sugar, you could add an is_const() that simiply creats a new ConstPattern, but I don't think it's strictly necessary.

I could do that. How about is_tuple and is_tuple_get_item?

For my own edification, what did you use to auto-format the test file? I tried autopep8 on it a while back and got really bad results, so I ignored it after that (as does the linter).

I used yapf with the .style.yapf I put under the TVM home.

[mbaret]

Thanks for the PR so soon! Is there an example of how partition works on a constant match? In particular, does the constant remain propagated into the body?

On a general point, it's helpful if reformats are kept separate from feature additions for reviews so it's easy to see what the changes are.

[mbrookhart]

@mbaret Originally the partition pass embedded constants in the function body, but @comaniac filed #5662, and I responded with #5663, so it currently will lift the constants to the function arguments.

Do you prefer embedded constants?

[comaniac]

Thanks for the PR so soon! Is there an example of how partition works on a constant match? In particular, does the constant remain propagated into the body?

An example can be found in the unit test: https://github.com/apache/incubator-tvm/pull/5689/files#diff-f9920485e5e341787129348ce1985db9R213

Also, could you provide an example to show your expection of constant propogation? I've checked that all merge composite unit tests are passed with the current implement even for the DNNL codegen unit tests that use composite functions with parameter binding. It would be better if you could include all your use cases in the unit tests so that we could on the same page.

On a general point, it's helpful if reformats are kept separate from feature additions for reviews so it's easy to see what the changes are.

Yeah the bugfix should be put to a separate PR for sure. For re-formatting, I didn't intentionally reformat the unit test file. VSCode by default auto-formats the file once I paste some code. I was even considering to revert auto-formatting due to the reason you pointed out. However, I checked the re-formatting and it is mostly just blank and too-long lines, so I think it should be fine for this small feature PR.

[mbaret]

%2 = fn (%input: Tensor[(1, 224, 224, 3), uint8], Composite="qnn_conv2d") -> Tensor[(1, 224, 224, 64), uint8] {
    %0 = qnn.conv2d(%input, meta[relay.Constant][0] /* ty=Tensor[(3, 3, 3, 64), uint8] */ /* ty=Tensor[(3, 3, 3, 64), uint8] */, 115 /* ty=int32 */, 134 /* ty=int32 */, 1f /* ty=float32 */, 0.00503911f /* ty=float32 */, padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO", out_dtype="int32") /* ty=Tensor[(1, 224, 224, 64), int32] */;
    %1 = nn.bias_add(%0, meta[relay.Constant][1] /* ty=Tensor[(64), int32] */ /* ty=Tensor[(64), int32] */, axis=3) /* ty=Tensor[(1, 224, 224, 64), int32] */;
    qnn.requantize(%1, 0.00503911f /* ty=float32 */, 0 /* ty=int32 */, 3.26957f /* ty=float32 */, 0 /* ty=int32 */, out_dtype="uint8") /* ty=Tensor[(1, 224, 224, 64), uint8] */
  };

Would be an example of the composite function I'd expect if I matched the weight of the conv as a constant (this is the old merge composite behaviour for that case).

I think the most logical behaviour is for constants to remain embedded where they are explicitly part of the pattern and to lift them where they are not.

[masahi]

I might be missing something, but I expect constants to be available to external codegen at compile time. Otherwise we cannot do constant folding. Sorry I didn't put thought on this issue when merging #5663

@comaniac Under the current implementation of merge composite + pattern partitioning, do you hit this visitor https://github.com/apache/incubator-tvm/blob/master/src/relay/backend/contrib/dnnl/codegen.cc#L156 during DNNL codegen?

[mbrookhart]

☝️ I'll let you guys discuss the appropriate behavior, there's seems to be some complication to this.

In the mean time, this was my initial assumption:

I expect constants to be available to external codegen at compile time. Otherwise we cannot do constant folding.

But it doesn't match the fusion pass.

I think the most logical behaviour is for constants to remain embedded where they are explicitly part of the pattern and to lift them where they are not.

I think about how to do that.

[comaniac]

I made an example and tested it with old MergeComposite pass (master commit 6100112). It behaves as I expected. The partitioned composite function still has 3 arguments even %w has been bind. @mbaret could you double check?

import tvm
from tvm import relay
from tvm import tir
from tvm.relay.testing import run_opt_pass

from tvm.relay.build_module import bind_params_by_name
import numpy as np


# Make a graph
x = relay.var('x', shape=(1, 3, 224, 224))
w = relay.var('w', shape=(3, 3, 3, 3))
b = relay.var('b', shape=(3,))

conv2d = relay.op.nn.conv2d(x, w)
out = relay.op.nn.bias_add(conv2d, b)
func = relay.Function([x, w, b], out)
mod = tvm.IRModule.from_expr(func)

mod["main"] = bind_params_by_name(mod["main"],
                                  {'w': tvm.nd.array(np.ones(shape=(3, 3, 3, 3)))})
print('=== Before ===')
print(mod['main'].body)


def pat():
    x = relay.var('x', shape=(1, 3, 224, 224))
    w = relay.var('w', shape=(3, 3, 3, 3))
    b = relay.var('b', shape=(3,))

    conv2d = relay.op.nn.conv2d(x, w)
    out = relay.op.nn.bias_add(conv2d, b)
    return out

pattern_table = [('pat', pat())]
result = run_opt_pass(mod['main'], relay.transform.MergeComposite(pattern_table))
print('=== After ===')
print(result)

=== Before ===
free_var %x: Tensor[(1, 3, 224, 224), float32]
%0 = nn.conv2d(%x, meta[relay.Constant][0] /* ty=Tensor[(3, 3, 3, 3), float64] */ /* ty=Tensor[(3, 3, 3, 3), float64] */, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 3, 222, 222), float32] */;
free_var %b: Tensor[(3), float32]
nn.bias_add(%0, %b) /* ty=Tensor[(1, 3, 222, 222), float32] */
// meta data omitted. you can use show_meta_data=True to include meta data
=== After ===
fn (%x: Tensor[(1, 3, 224, 224), float32], %b: Tensor[(3), float32]) -> Tensor[(1, 3, 222, 222), float32] {
  %1 = fn (%x1: Tensor[(1, 3, 224, 224), float32], %w: Tensor[(3, 3, 3, 3), float64], %b1: Tensor[(3), float32], Composite="pat") -> Tensor[(1, 3, 222, 222), float32] {
    %0 = nn.conv2d(%x1, %w, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 3, 222, 222), float32] */;
    nn.bias_add(%0, %b1) /* ty=Tensor[(1, 3, 222, 222), float32] */
  };
  %1(%x, meta[relay.Constant][0] /* ty=Tensor[(3, 3, 3, 3), float64] */ /* ty=Tensor[(3, 3, 3, 3), float64] */, %b) /* ty=Tensor[(1, 3, 222, 222), float32] */
}
// meta data omitted. you can use show_meta_data=True to include meta data

@masahi I just checked with this PR. The unit test hits the line you pointed out twice.

[mbaret]

Try changing w in the pattern to a relay.const rather than a var.

[masahi]

@masahi I just checked with this PR. The unit test hits the line you pointed out twice.

hmm something seems off to me. On the dnnl fused mobilenet tests, I think all params in the network should go through that visitor. In #5310, @zhiics improved how to handle big constants, which made running fused mobilenet possible. If the visitor is called only twice, we shouldn't have had such issues at all.

[comaniac]

OK I can finally reproduce this case:

fn (%x: Tensor[(1, 3, 224, 224), float32], %b: Tensor[(3), float32]) -> Tensor[(1, 3, 222, 222), float32] {
  %1 = fn (%x1: Tensor[(1, 3, 224, 224), float32], %b1: Tensor[(3), float32], Composite="pat") -> Tensor[(1, 3, 222, 222), float32] {
    %0 = nn.conv2d(%x1, meta[relay.Constant][0] /* ty=Tensor[(3, 3, 3, 3), float64] */ /* ty=Tensor[(3, 3, 3, 3), float64] */, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 3, 222, 222), float32] */;
    nn.bias_add(%0, %b1) /* ty=Tensor[(1, 3, 222, 222), float32] */
  };
  %1(%x, %b) /* ty=Tensor[(1, 3, 222, 222), float32] */
}
// meta data omitted. you can use show_meta_data=True to include meta data

Combining wih the previous example, we know that old MergeComposite preserves constant node when the pattern is also a constant; otherwise it will lift constant nodes. In the DNNL codegen, the patterns are all specified as VarNode, so the case that matches a pattern with constant nodess never shows up in unit tests and that's why we missed it.

However, I personally think this behavior is weird. General speaking, whatever we specify var or const in the pattern, we may need or not need constant lifting. It seems to me that this partition behavior should not be bind with patterns but should be a partition option.

Anyways, I think the behavior of constant lifting should be discussed in another topic and does not relate to this PR.

[mbaret]

It's not really the case that old merge composite explicitly 'lifts' constant nodes if they're vars, in old merge composite vars indicate inputs to the pattern (so they don't need to explicitly match against a VarNode).

My expectation would be if the pattern requires there to be a constant node, I should see that constant node in the partitioned function. In other words, the body of the function should still match the pattern used to create it.

[mbrookhart]

I think I agree with @mbaret here. I think I see a simple way to do that behavior by default, why don't I post another PR in that direction after this goes in?

[comaniac]

That's the behavior I just realized with above examples. In this case, what would be the behavior if we specify VarPattern('x') | ConstantPattern()?

[mbrookhart]

I think it depends on which one matched? If it's a var, we lift it, if it's a constant, we embed?

[comaniac]

OK so here is a conclusion for another follow-up PR to deal with constant lifting.

• Pattern input ConstantPattern:
  • Only match constant nodes and keep them in the partitioned functions. In this case, the arguments of partitioned functions will be reduced.
• Pattern input VarPattern('x'):
  • Match only the var node with a specified name hint. In this case, the arguments of partitioned function is fixed.
• Pattern input wildcard or VarPattern('x') | ConstantPattern:
  • Match both constant and var node but life constant nodes. In this case, the arguments of partitioned function is fixed whever it matches constant or var.

btw, IMHO, for clearer semantic, it might be better if we support VarPattern() that matches a var node with any name hints but not constant node.

[mbrookhart]

I'd liked to refine slightly:

• Pattern input ExprPattern(relay.const):
  • Only match constant nodes that match the value and embed them in the partitioned functions. In this case, the arguments of partitioned functions will be reduced.
• Pattern input ConstantPattern:
  • Only match constant nodes and embed them in the partitioned functions. In this case, the arguments of partitioned functions will be reduced.
• Pattern input VarPattern('x'):
  • Match only the var node with an optional specified name hint. In this case, the arguments of partitioned function is fixed.
• Pattern input wildcard:
  • Match anything. In this case, the arguments of partitioned function is fixed.
• Pattern AltPattern:
  • Match either lhs or rhs. depending on which side matched, recursively apply this matching logic and embed constants appropriately based on the contained pattern.

[mbrookhart]

@comaniac The input VarPattern name is already optional, if you use is_input() it will match any VarNode
https://github.com/apache/incubator-tvm/blob/a072da0588c542757d2815832b7f010f530b2428/src/relay/ir/dataflow_matcher.cc#L386-L395
https://github.com/apache/incubator-tvm/blob/a072da0588c542757d2815832b7f010f530b2428/python/tvm/relay/dataflow_pattern/__init__.py#L171-L185

I'm not sure that made it into the document or test cases, I'll make sure to update in the followup PR.

[masahi]

bound, can be fixed in the next PR

[masahi]

@comaniac @mbrookhart good to go?

[comaniac]

@comaniac @mbrookhart good to go?

Yeah I'm fine with that.