anijain2305 commented on a change in pull request #4771: [Relay] Added Merge Composite pass URL: https://github.com/apache/incubator-tvm/pull/4771#discussion_r376750916
########## File path: tests/python/relay/test_pass_merge_composite.py ########## @@ -0,0 +1,439 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Unit tests for merge composite.""" +from tvm import expr +from tvm import relay +from tvm.relay.testing import run_opt_pass + +""" +The merge composite pass is designed to merge multiple relay operators, that +match a given pattern, and combine them into a single relay function. + +For example suppose we have the graph: + + conv2d + | (merge composite pass) + bias_add ====> conv2d_bias_relu + | (our target) + relu + +Our Relay IR before the pass: + fn (%data: Tensor[(1, 512, 28, 28), float32], %kernel: Tensor[(256, 512, 1, 1), float32], + %bias: Tensor[(256), float32]) -> Tensor[(1, 256, 28, 28), float32] { + %0 = nn.conv2d(%data, %kernel, kernel_size=[1, 1]) + /* ty=Tensor[(1, 256, 28, 28), float32] */; + %1 = nn.bias_add(%0, %bias) /* ty=Tensor[(1, 256, 28, 28), float32] */; + nn.relu(%1) /* ty=Tensor[(1, 256, 28, 28), float32] */ + } + +Our Relay IR after the pass: + fn (%data: Tensor[(1, 512, 28, 28), float32], %kernel: Tensor[(256, 512, 1, 1), float32], + %bias: Tensor[(256), float32]) -> Tensor[(1, 256, 28, 28), float32] { + %2 = fn (%x: Tensor[(1, 512, 28, 28), float32], %y: Tensor[(256, 512, 1, 1), float32], + %z: Tensor[(256), float32], Primitive=1, Composite="conv2d_bias_relu") -> + Tensor[(1, 256, 28, 28), float32] { + %0 = nn.conv2d(%x, %y, kernel_size=[1, 1]) /* ty=Tensor[(1, 256, 28, 28), float32] */; + %1 = nn.bias_add(%0, %z) /* ty=Tensor[(1, 256, 28, 28), float32] */; + nn.relu(%1) /* ty=Tensor[(1, 256, 28, 28), float32] */ + }; + %2(%data, %kernel, %bias) /* ty=Tensor[(1, 256, 28, 28), float32] */ + } + +As you can see in the second relay example, the pattern we specified has been wrapped +in a function. The function is then called, producing the same result as the first relay +example. + +One convenient use for this pass is to offload multiple operators to a single external +codegen function. +""" + + +def make_add_sub_mul_pattern(): + """Create a pattern to match the following graph. + + add sub + \ / + \ / + mul + """ + x = relay.var('x') + y = relay.var('y') + add_node = relay.add(x, y) + sub_node = relay.subtract(x, y) + mul_node = relay.multiply(add_node, sub_node) + return mul_node + + +def make_add_relu_pattern(): + """Create a pattern to match the following graph. + + add + | + relu + """ + x = relay.var('x') + y = relay.var('y') + add_node = relay.add(x, y) + r = relay.nn.relu(add_node) + return r + + +def make_conv_bias_relu_pattern(): + """Create a pattern to match the following graph. + + conv2d + | + bias_add + | + relu + """ + x = relay.var('x') + y = relay.var('y') + z = relay.var('z') + conv_node = relay.nn.conv2d(x, y) + bias_node = relay.nn.bias_add(conv_node, z) + r = relay.nn.relu(bias_node) + return r + + +def test_simple_merge(): + """Test composite function is correctly produced from simple graph. + + We could expect the pattern `make_add_relu_pattern` to be merged + into a single op `add_relu`. + + a b + \ / a b + add ====> \ / + | add_relu + relu + + """ + pattern_table = [ + ("add_relu", make_add_relu_pattern()) + ] + + def before(): + a = relay.var('a', shape=(10, 10)) + b = relay.var('b', shape=(10, 10)) + add_node = relay.add(a, b) + r = relay.nn.relu(add_node) + return relay.Function([a, b], r) + + def expected(): + a = relay.var('a', shape=(10, 10)) + b = relay.var('b', shape=(10, 10)) + + # add_relu function + in_1 = relay.var('in_1', shape=(10, 10)) + in_2 = relay.var('in_2', shape=(10, 10)) + add_node = relay.add(in_1, in_2) + relu_node = relay.nn.relu(add_node) + add_relu = relay.Function([in_1, in_2], relu_node) + + # merged function + r = relay.Call(add_relu, [a, b]) + return relay.Function([a, b], r) + + result = run_opt_pass(before(), relay.transform.MergeComposite(pattern_table)) + assert not relay.analysis.free_vars(result) + expected = run_opt_pass(expected(), relay.transform.InferType()) + assert relay.analysis.alpha_equal(result, expected) + + +def test_branch_merge(): + """Test composite function is correctly produced from branching graph. + + We would expect the pattern `make_add_sub_mul_pattern` to be merged + into a single op `add_sub_mul`. + + a b a b + \/ \/ + add sub a b + \ / \/ Review comment: Can we add 1 more testcase where Var('a') = Var('c') in the primary graph? Maybe, my previous comment already covers that. ---------------------------------------------------------------- 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 With regards, Apache Git Services
