Registering operator attributes for Relay pattern matching

StMa · January 18, 2021, 1:42am

Hello,

This question is regarding the pattern matching in Relay.

Information from the following page was used as a reference: https://tvm.apache.org/docs/langref/relay_pattern.html

For a pattern like this:

self.x = wildcard()
mean_op = is_op('mean')(self.x)
self.pattern = \
               mean_op.has_attr({"axis":        3 }) | \
               mean_op.has_attr({"axis":    (1, 3)}) | \
               mean_op.has_attr({"axis":    (2, 3)}) | \
               mean_op.has_attr({"axis": (1, 2, 3)})
self.pattern = self.pattern(self.x)

I was trying to match the following operator

relay.mean(data, axis=axis, keepdims=keepdims, exclude=exclude)

for the axis values outlined in the pattern above.

This didn’t work at first since the ‘axis’ attribute needed to be registered apparently.

The mean operator with the axis set to 3 could be matched by the pattern above after using

tvm.ir.register_op_attr("mean", "axis", value=axis)

but could not be matched for other values in the pattern, i.e. the tuple values for the axis.

The question would be: is the pattern written correctly and if it is, how can values that are tuples (in this case) be matched as the values for the axis attribute?

I noticed that the value type that register_op_attr returns is tvm.ir.container.Array after registering the axis attribute with a value that is a tuple, ex. (1, 2, 3); could that be a problem, since the value that will be matched in the operator is a python tuple?

Thanks in advance.

lixiaoquan · January 18, 2021, 5:37am

Current pattern matching code doesn’t handle Array in attrs

github.com

apache/tvm/blob/052ad3d92d20abdf221600005c2ccb130e39b6b4/src/relay/ir/dataflow_matcher.cc#L109


      ClearMap(watermark);
    }
    return out;
  }
}

bool DFPatternMatcher::VisitDFPattern_(const AltPatternNode* op, const Expr& expr) {
  return VisitDFPattern(op->left, expr) || VisitDFPattern(op->right, expr);
}

bool MatchRetValue(const ObjectRef& lhs, const TVMRetValue& rhs) {
  switch (rhs.type_code()) {
    case kDLInt:
      if (auto* val = lhs.as<IntImmNode>()) {
        return val->value == rhs.operator int64_t();
      }
      break;
    case kDLFloat:
      if (auto* val = lhs.as<FloatImmNode>()) {
        return val->value == rhs.operator double();
      }

comaniac · January 18, 2021, 6:03pm

It’s right. The semantic of matching arrays is relative ambiguous. If you really need to do so, you can consider either of the following, depending on the action you want to take after matching a pattern:

If you attempt to rewrite the pattern, add your person check to the rewrite function and simply return the original graph if not matched.

If you attempt to partition the pattern, customize the partition check callback to reject the patterns that do not match your axis.