[GitHub] [incubator-tvm] mbaret commented on a change in pull request #5646: Add a check Callback to the Pattern Paritioner

[GitBox]

mbaret commented on a change in pull request #5646:
URL: https://github.com/apache/incubator-tvm/pull/5646#discussion_r428940644



##########
File path: include/tvm/relay/dataflow_matcher.h
##########
@@ -90,7 +91,8 @@ Expr RewritePatterns(Array<DFPatternCallback> callbacks, Expr 
expr);
  *
  * \return Return the paritioned Expr.
  */
-Expr PartitionPattern(DFPattern pattern, Expr expr);

Review comment:
       The history's a bit confusing here, did the last version not have an 
attrs argument?

##########
File path: src/relay/ir/dataflow_matcher.cc
##########
@@ -728,16 +729,17 @@ class PatternPartitioner : protected MixedModeMutator {
   Map<std::string, ObjectRef> attrs_;
   std::vector<PatternGrouper::Group> groups_;
   std::unordered_map<Expr, int, ObjectHash, ObjectEqual> gid_assignments_;
+  PackedFunc check_;
 };
 
-Expr PartitionPattern(DFPattern pattern, Expr expr, Map<std::string, 
ObjectRef> attrs) {
-  return PatternPartitioner().Partition(pattern, expr, attrs);
+Expr PartitionPattern(DFPattern pattern, Expr expr, Map<std::string, 
ObjectRef> attrs,
+                      PackedFunc check) {
+  return PatternPartitioner().Partition(pattern, expr, attrs, check);
 }
 
 TVM_REGISTER_GLOBAL("relay.dataflow_pattern.partition")
-    .set_body_typed([](DFPattern pattern, Expr expr, Map<std::string, 
ObjectRef> attrs) {
-      return PartitionPattern(pattern, expr, attrs);
-    });
+    .set_body_typed([](DFPattern pattern, Expr expr, Map<std::string, 
ObjectRef> attrs,
+                       PackedFunc check) { return PartitionPattern(pattern, 
expr, attrs, check); });

Review comment:
       Probably keep the original formatting here? Not sure what the style 
convention is though but it definitely reads more clearly.

##########
File path: tests/python/relay/test_dataflow_pattern.py
##########
@@ -959,4 +984,5 @@ def test_parition_double_batchnorm():
     test_quadruple_partition_dominator()
     test_parition_batchnorm()
     test_parition_double_batchnorm()
+    test_parition_check()

Review comment:
       test_partition




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org