[QNN][Relay] Calling Dialect passes from inside Relay Build API.

Author: Ubuntu

include/tvm/relay/qnn/transform.h

@@ -0,0 +1,60 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relay/qnn/transform.h
+ *
+ * This file implements a pass manager for QNN ops using Relay Pass manager.
+ */
+#ifndef TVM_RELAY_QNN_TRANSFORM_H_
+#define TVM_RELAY_QNN_TRANSFORM_H_
+
+#include <tvm/relay/transform.h>
+
+namespace tvm {
+namespace relay {
+
+using relay::transform::Pass;
+
+namespace qnn {
+namespace transform {
+
+/*!
+ * \brief Legalizes a QNN expr.
+ * \param Contains specifically two types of Legalizations. First, converts/Lowers an expression
+ * containing QNN ops to an expression containing only core Relay ops. Each QNN op is lowered to a
+ * sequence of exisiting Relay ops. This is a target-independent pass. One can register the
+ * lowering/transformation function for this op using FTVMQnnCanonicalize attr_name for FTVMLegalize
+ * op attribute. Second, as opposed to Relay Legalize, this one legalizes only QNN ops. One can
+ * register a transformation/legalization function for an op by using the FTVMQnnLegalize attr_name
+ * for FTVMLegalize op attribute. The isolation of QNN and Relay Legalize gives us separation of
+ * concerns, leading to a better software practice. The legalization can be configured to happen per
+ * target.
+ *
+ * \return The pass.
+ */
+Pass Legalize();
+
+}  // namespace transform
+
+}  // namespace qnn
+}  // namespace relay
+}  // namespace tvm
+
+#endif  // TVM_RELAY_QNN_TRANSFORM_H_

src/relay/backend/build_module.cc

@@ -27,6 +27,7 @@
 #include <tvm/runtime/vm.h>
 #include <tvm/relay/expr.h>
 #include <tvm/relay/transform.h>
+#include <tvm/relay/qnn/transform.h>
 #include <memory>
 
 #include "utils.h"
@@ -282,6 +283,15 @@ class RelayBuildModule : public runtime::ModuleNode {
       const TargetsMap& targets,
       const std::unordered_map<std::string, runtime::NDArray>& params) {
     Array<Pass> pass_seqs;
+
+    // Run all dialect legalization passes.
+    pass_seqs.push_back(relay::qnn::transform::Legalize());
+
+    // Legalize pass is restricted to homogeneous execution for now.
+    if (targets.size() == 1) {
+      pass_seqs.push_back(transform::Legalize());
+    }
+
     pass_seqs.push_back(transform::SimplifyInference());
     PackedFunc fskip = PackedFunc([](TVMArgs args, TVMRetValue* rv) {
       Expr expr = args[0];
@@ -304,11 +314,6 @@ class RelayBuildModule : public runtime::ModuleNode {
     pass_seqs.push_back(transform::CanonicalizeCast());
     pass_seqs.push_back(transform::CanonicalizeOps());
 
-    // Legalize pass is restricted to homogeneous execution for now.
-    if (targets.size() == 1) {
-      pass_seqs.push_back(transform::Legalize());
-    }
-
     // Alter layout transformation is only applied to homogeneous execution yet.
     if (targets.size() == 1) {
       pass_seqs.push_back(transform::AlterOpLayout());

src/relay/pass/legalize.cc

@@ -95,7 +95,7 @@ Pass Legalize(const std::string& legalize_map_attr_name) {
       [=](Function f, Module m, PassContext pc) {
         return Downcast<Function>(relay::legalize::Legalize(f, legalize_map_attr_name));
       };
-  return CreateFunctionPass(pass_func, 3, "Legalize", {ir::StringImm::make("InferType")});
+  return CreateFunctionPass(pass_func, 0, "Legalize", {ir::StringImm::make("InferType")});
 }
 
 TVM_REGISTER_API("relay._transform.Legalize").set_body_typed(Legalize);

src/relay/qnn/pass/legalize.cc

@@ -0,0 +1,47 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file relay/qnn/pass/legalize.cc
+ * \brief The Legalize wrapper for QNN.
+ */
+
+#include <tvm/relay/qnn/transform.h>
+
+namespace tvm {
+namespace relay {
+namespace qnn {
+
+namespace transform {
+
+Pass Legalize() {
+  Array<Pass> pass_seqs;
+  pass_seqs.push_back(relay::transform::Legalize("FTVMQnnLegalize"));
+  pass_seqs.push_back(relay::transform::Legalize("FTVMQnnCanonicalize"));
+  relay::transform::Pass seq = relay::transform::Sequential(pass_seqs);
+  return seq;
+}
+
+TVM_REGISTER_API("relay.qnn._transform.Legalize").set_body_typed(Legalize);
+
+}  // namespace transform
+
+}  // namespace qnn
+}  // namespace relay
+}  // namespace tvm

tests/python/relay/test_op_qnn_conv2d.py

@@ -78,7 +78,6 @@ def get_qnn_func(data,
 
     mod = relay.Function(relay.analysis.free_vars(func), func)
     mod = relay.Module.from_expr(mod)
-    mod = relay.qnn.transform.CanonicalizeOps()(mod)
     return mod
 
 def get_funcs(data_shape,

tests/python/relay/test_op_qnn_dequantize.py

@@ -31,7 +31,6 @@ def quantize_test_driver(in_dtype, quant_args, in_data, verify_output_data):
                                                    input_zero_point=input_zero_point)
         mod = relay.Function(relay.analysis.free_vars(quantized_output), quantized_output)
         mod = relay.Module.from_expr(mod)
-        mod = relay.qnn.transform.CanonicalizeOps()(mod)
         with relay.build_config(opt_level=3):
             graph, lib, params = relay.build(mod, "llvm", params=None)
             rt_mod = graph_runtime.create(graph, lib, ctx=tvm.cpu(0))

tests/python/relay/test_op_qnn_quantize.py

@@ -31,7 +31,6 @@ def quantize_test_driver(in_dtype, quant_args, out_dtype, in_data, verify_output
                                                  output_zero_point=output_zero_point,out_dtype=out_dtype)
         mod = relay.Function(relay.analysis.free_vars(quantized_output), quantized_output)
         mod = relay.Module.from_expr(mod)
-        mod = relay.qnn.transform.CanonicalizeOps()(mod)
         with relay.build_config(opt_level=3):
             graph, lib, params = relay.build(mod, "llvm", params=None)
             rt_mod = graph_runtime.create(graph, lib, ctx=tvm.cpu(0))

tests/python/relay/test_op_qnn_requantize.py

@@ -49,7 +49,6 @@ def get_mod(data_shape, data_dtype, out_dtype, input_scale, output_scale,
 
         mod = relay.Function(relay.analysis.free_vars(mod), mod)
         mod = relay.Module.from_expr(mod)
-        mod = relay.qnn.transform.CanonicalizeOps()(mod)
         return mod
 
     def same_scale_test():