[Relay] Expose qnn ops directly from relay.qnn module

python/tvm/relay/qnn/__init__.py

@@ -14,8 +14,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=wildcard-import
+# pylint: disable=wildcard-import,redefined-builtin
 """QNN dialect operators and IR passes."""
 from __future__ import absolute_import as _abs
 from . import op
 from . import transform
+from .op.qnn import *

tests/python/contrib/test_hexagon/test_qnn_op_integration.py

@@ -35,8 +35,8 @@
 def test_disable_qnn_legalize_pass():
     """No QNN pass test."""
     x = relay.var("x", shape=(4, 8), dtype="float32")
-    op0 = relay.qnn.op.quantize(x, relay.const(2.0), relay.const(10), out_dtype="uint8")
-    op1 = relay.qnn.op.dequantize(op0, relay.const(0.5), relay.const(5))
+    op0 = relay.qnn.quantize(x, relay.const(2.0), relay.const(10), out_dtype="uint8")
+    op1 = relay.qnn.dequantize(op0, relay.const(0.5), relay.const(5))
     relay_mod = tvm.IRModule.from_expr(op1)
 
     target_hexagon = tvm.target.hexagon("v68")
@@ -133,7 +133,7 @@ def test_qnn_quantize(self, hexagon_session: Session, odtype, input_shape):
 
         def gen_relay_expr_qnn(output_scale, output_zero_point):
             data = relay.var("data", shape=input_shape, dtype="float32")
-            qnn_quantize = relay.qnn.op.quantize(
+            qnn_quantize = relay.qnn.quantize(
                 data,
                 output_scale=relay.const(output_scale),
                 output_zero_point=relay.const(output_zero_point),
@@ -167,7 +167,7 @@ def test_qnn_dequantize(self, hexagon_session: Session, idtype, input_shape):
 
         def gen_relay_expr_qnn(dtype, input_scale, input_zero_point):
             data = relay.var("data", shape=input_shape, dtype=dtype)
-            qnn_dequantize = relay.qnn.op.dequantize(
+            qnn_dequantize = relay.qnn.dequantize(
                 data,
                 input_scale=relay.const(input_scale),
                 input_zero_point=relay.const(input_zero_point),
@@ -199,7 +199,7 @@ def test_qnn_requantize(self, hexagon_session: Session):
         data_shape = [256]
         data = relay.var("data", shape=data_shape, dtype="int32")
 
-        op = relay.qnn.op.requantize(
+        op = relay.qnn.requantize(
             data,
             input_scale=relay.const(0.156),
             input_zero_point=relay.const(2),
@@ -256,7 +256,7 @@ def gen_relay_expr_qnn(
             dtype, input_scale, input_zero_point, output_scale, output_zero_point
         ):
             data = relay.var("data", shape=input_shape, dtype=dtype)
-            qnn_avg_pool = relay.qnn.op.avg_pool2d(
+            qnn_avg_pool = relay.qnn.avg_pool2d(
                 data,
                 input_scale=relay.const(input_scale),
                 input_zero_point=relay.const(input_zero_point),
@@ -305,7 +305,7 @@ def gen_relay_expr_qnn(
 class TestQnnBinaryOp:
     """QNN binary op test class"""
 
-    operation = tvm.testing.parameter(relay.qnn.op.add, relay.qnn.op.subtract, relay.qnn.op.mul)
+    operation = tvm.testing.parameter(relay.qnn.add, relay.qnn.subtract, relay.qnn.mul)
     dtype = tvm.testing.parameter("uint8", "int8")
     input_shape = tvm.testing.parameter([256], [4, 256])
 
@@ -375,7 +375,7 @@ def test_qnn_concatenate(self, hexagon_session: Session):
         input_y = relay.var("y", shape=y_shape, dtype="uint8")
         input_z = relay.var("z", shape=z_shape, dtype="uint8")
 
-        op = relay.qnn.op.concatenate(
+        op = relay.qnn.concatenate(
             (input_x, input_y, input_z),
             input_scales=(relay.const(0.3), relay.const(0.7), relay.const(1.3)),
             input_zero_points=(relay.const(0), relay.const(1), relay.const(2)),
@@ -404,9 +404,9 @@ def test_qnn_quantize_conv2d_requantize(self, hexagon_session: Session):
         weight_shape = [16, 8, 3, 3]
         data = relay.var("data", shape=data_shape, dtype="float32")
         weight = relay.var("weight", shape=weight_shape, dtype="float32")
-        op0 = relay.qnn.op.quantize(data, relay.const(0.078), relay.const(0), out_dtype="uint8")
-        op1 = relay.qnn.op.quantize(weight, relay.const(0.07), relay.const(0), out_dtype="int8")
-        op2 = relay.qnn.op.conv2d(
+        op0 = relay.qnn.quantize(data, relay.const(0.078), relay.const(0), out_dtype="uint8")
+        op1 = relay.qnn.quantize(weight, relay.const(0.07), relay.const(0), out_dtype="int8")
+        op2 = relay.qnn.conv2d(
             op0,
             op1,
             input_zero_point=relay.const(0),
@@ -417,7 +417,7 @@ def test_qnn_quantize_conv2d_requantize(self, hexagon_session: Session):
             channels=16,
             kernel_size=[3, 3],
         )
-        op5 = relay.qnn.op.requantize(
+        op5 = relay.qnn.requantize(
             op2,
             input_scale=relay.const(0.05),
             input_zero_point=relay.const(0),
@@ -448,11 +448,11 @@ def test_alter_layout_qnn_dense(self):
         wscale = relay.const(0.37)
 
         def before():
-            return relay.qnn.op.dense(data, weight, zero, zero, iscale, wscale, units=None)
+            return relay.qnn.dense(data, weight, zero, zero, iscale, wscale, units=None)
 
         def expected():
             op0 = relay.layout_transform(weight, src_layout="NC", dst_layout="NC32n4c")
-            return relay.qnn.op.contrib_dense_pack(data, op0, zero, zero, iscale, wscale, "NC32n4c")
+            return relay.qnn.contrib_dense_pack(data, op0, zero, zero, iscale, wscale, "NC32n4c")
 
         target = tvm.target.hexagon("v68")
         with tvm.target.Target(target):
@@ -478,7 +478,7 @@ def test_qnn_dense_biasadd_requantize(self, hexagon_session: Session, dtype, n_d
         weight = relay.var("weight", shape=weight_shape, dtype=dtype)
         bias = relay.var("bias", shape=bias_shape, dtype="int32")
 
-        op0 = relay.qnn.op.dense(
+        op0 = relay.qnn.dense(
             data,
             weight,
             input_zero_point=relay.const(2),
@@ -488,7 +488,7 @@ def test_qnn_dense_biasadd_requantize(self, hexagon_session: Session, dtype, n_d
             units=None,
         )
         op1 = relay.nn.bias_add(op0, bias)
-        op2 = relay.qnn.op.requantize(
+        op2 = relay.qnn.requantize(
             op1,
             input_scale=relay.const(1.3),
             input_zero_point=relay.const(4),
@@ -520,7 +520,7 @@ def test_qnn_dense_requantize(self, hexagon_session: Session):
         data = relay.var("data", shape=data_shape, dtype="uint8")
         weight = relay.var("weight", shape=weight_shape, dtype="int8")
 
-        op0 = relay.qnn.op.dense(
+        op0 = relay.qnn.dense(
             data,
             weight,
             input_zero_point=relay.const(0),
@@ -529,7 +529,7 @@ def test_qnn_dense_requantize(self, hexagon_session: Session):
             kernel_scale=relay.const(0.19),
             units=64,
         )
-        op1 = relay.qnn.op.requantize(
+        op1 = relay.qnn.requantize(
             op0,
             input_scale=relay.const(0.1),
             input_zero_point=relay.const(0),
@@ -558,7 +558,7 @@ def test_qnn_tanh(self, hexagon_session: Session):
         data_shape = [256]
         data = relay.var("data", shape=data_shape, dtype="uint8")
 
-        op = relay.qnn.op.tanh(
+        op = relay.qnn.tanh(
             data,
             scale=relay.const(0.518),
             zero_point=relay.const(137),

tests/python/contrib/test_hexagon/test_relay_simplify_qnn_concat.py

@@ -43,11 +43,11 @@ def get_test_module():
         ceil_mode=False,
         layout="NHWC",
     )
-    r2 = relay.qnn.op.requantize(q2, s2, z1, s5, z1, axis=1, out_dtype="uint8")
+    r2 = relay.qnn.requantize(q2, s2, z1, s5, z1, axis=1, out_dtype="uint8")
     q_tuple = relay.expr.Tuple([r1, r2, q3])
     s_tuple = relay.expr.Tuple([s4, s5, s3])
     z_tuple = relay.expr.Tuple([z1, z1, z1])
-    graph = relay.qnn.op.concatenate(q_tuple, s_tuple, z_tuple, s3, z1, axis=1)
+    graph = relay.qnn.concatenate(q_tuple, s_tuple, z_tuple, s3, z1, axis=1)
 
     func = relay.Function(relay.analysis.free_vars(graph), graph)
     mod = tvm.IRModule.from_expr(func)
@@ -72,12 +72,10 @@ def get_expected_output_module():
         ceil_mode=False,
         layout="NHWC",
     )
-    out_r1 = relay.qnn.op.requantize(
+    out_r1 = relay.qnn.requantize(
         nn_max_pool, out_s4, out_z1, out_s3, out_z1, axis=1, out_dtype="uint8"
     )
-    out_r2 = relay.qnn.op.requantize(
-        out_q2, out_s2, out_z1, out_s3, out_z1, axis=1, out_dtype="uint8"
-    )
+    out_r2 = relay.qnn.requantize(out_q2, out_s2, out_z1, out_s3, out_z1, axis=1, out_dtype="uint8")
     out_q_tuple = relay.expr.Tuple([out_r1, out_r2, out_q3])
     out_graph = relay.op.concatenate(out_q_tuple, axis=1)
 

tests/python/relay/qnn/test_clip_legalization.py

@@ -52,7 +52,7 @@ def test_removes_redundant_requantize_clip_ops(dtype, min_val, max_val, is_redun
     the clip operator match the min and max values of the data type."""
 
     input_var = relay.var("input", shape=(1, 3, 3, 4), dtype="int32")
-    out = relay.qnn.op.requantize(
+    out = relay.qnn.requantize(
         input_var,
         tvm_const(np.float32(1.0)),
         tvm_const(np.int32(0)),

tests/python/relay/qnn/test_qnn_channel_stripping.py

@@ -58,7 +58,7 @@ def make_test_conv_depthwise_conv():
     input_scale_1 = np.float32(0.5)
     output_scale_1 = np.array([0.5, 2.0, 0.25, 4.0], dtype="float32")
 
-    out = relay.qnn.op.conv2d(
+    out = relay.qnn.conv2d(
         input_var,
         tvm_const(kernel_1),
         tvm_const(np.int32(-128)),
@@ -80,7 +80,7 @@ def make_test_conv_depthwise_conv():
     )
 
     input_scale_2 = np.float32(0.25)
-    out = relay.qnn.op.requantize(
+    out = relay.qnn.requantize(
         out,
         tvm_const(input_scale_1 * output_scale_1),
         tvm_const(np.int32(0)),
@@ -106,7 +106,7 @@ def make_test_conv_depthwise_conv():
         dtype="int8",
     ).reshape((3, 3, 4, 1))
     output_scale_2 = np.array([0.25, 0.125, 2.0, 0.125], dtype="float32")
-    out = relay.qnn.op.conv2d(
+    out = relay.qnn.conv2d(
         out,
         tvm_const(kernel_2),
         tvm_const(np.int32(-128)),
@@ -129,7 +129,7 @@ def make_test_conv_depthwise_conv():
     )
 
     input_scale_3 = np.float32(0.125)
-    out = relay.qnn.op.requantize(
+    out = relay.qnn.requantize(
         out,
         tvm_const(input_scale_2 * output_scale_2),
         tvm_const(np.int32(0)),
@@ -145,7 +145,7 @@ def make_test_conv_depthwise_conv():
     ).reshape((1, 1, 4, 4))
     output_scale_3 = np.array([0.25, 0.125, 1.0, 0.5], dtype="float32")
 
-    out = relay.qnn.op.conv2d(
+    out = relay.qnn.conv2d(
         out,
         tvm_const(kernel_3),
         tvm_const(np.int32(-128)),
@@ -181,7 +181,7 @@ def make_test_conv_pool_dense():
     input_scale = np.float32(0.029626124)
     output_scale = np.array([0.5, 2.0, 0.25, 4.0], dtype="float32")
 
-    out = relay.qnn.op.conv2d(
+    out = relay.qnn.conv2d(
         input_var,
         tvm_const(kernel),
         tvm_const(np.int32(-128)),
@@ -202,7 +202,7 @@ def make_test_conv_pool_dense():
         axis=3,
     )
 
-    out = relay.qnn.op.requantize(
+    out = relay.qnn.requantize(
         out,
         tvm_const(input_scale * output_scale),
         tvm_const(np.int32(0)),
@@ -226,7 +226,7 @@ def make_test_conv_pool_dense():
     out = relay.reshape(out, newshape=[-1, 4])
 
     dense_weights = np.array([[15, -2, -3, 11], [12, -10, 13, -10]], dtype="int8")
-    out = relay.qnn.op.dense(
+    out = relay.qnn.dense(
         out,
         tvm_const(dense_weights),
         tvm_const(np.int32(-128)),

tests/python/relay/test_op_qnn_add.py

@@ -25,7 +25,7 @@ def test_tflite_same_io_qnn_params():
 
     x = relay.var("x", shape=(1, 4), dtype=data_dtype)
     y = relay.var("y", shape=(1, 4), dtype=data_dtype)
-    z = relay.qnn.op.add(
+    z = relay.qnn.add(
         lhs=x,
         rhs=y,
         lhs_scale=relay.const(0.00784314, "float32"),
@@ -74,7 +74,7 @@ def test_tflite_different_io_qnn_params():
 
     x = relay.var("x", shape=(1, 4), dtype=data_dtype)
     y = relay.var("y", shape=(1, 4), dtype=data_dtype)
-    z = relay.qnn.op.add(
+    z = relay.qnn.add(
         lhs=x,
         rhs=y,
         lhs_scale=relay.const(0.0156863, "float32"),
@@ -123,7 +123,7 @@ def test_saturation():
     data_dtype = "uint8"
     x = relay.var("x", shape=(1, 4), dtype=data_dtype)
     y = relay.var("y", shape=(1, 4), dtype=data_dtype)
-    z = relay.qnn.op.add(
+    z = relay.qnn.add(
         lhs=x,
         rhs=y,
         lhs_scale=relay.const(0.125, "float32"),
@@ -151,7 +151,7 @@ def test_saturation():
     np.testing.assert_equal(op_res.numpy(), golden_output)
 
     # Same params, different scale
-    z = relay.qnn.op.add(
+    z = relay.qnn.add(
         lhs=x,
         rhs=y,
         lhs_scale=relay.const(0.125, "float32"),
@@ -178,7 +178,7 @@ def test_saturation():
     np.testing.assert_equal(op_res.numpy(), golden_output)
 
     # Same io params, different output scale
-    z = relay.qnn.op.add(
+    z = relay.qnn.add(
         lhs=x,
         rhs=y,
         lhs_scale=relay.const(0.125, "float32"),
@@ -205,7 +205,7 @@ def test_saturation():
     np.testing.assert_equal(op_res.numpy(), golden_output)
 
     # All params different
-    z = relay.qnn.op.add(
+    z = relay.qnn.add(
         lhs=x,
         rhs=y,
         lhs_scale=relay.const(0.5, "float32"),
@@ -237,7 +237,7 @@ def test_ignore_channel_axis():
 
     x = relay.var("x", shape=(4,), dtype=data_dtype)
     y = relay.var("y", shape=(4,), dtype=data_dtype)
-    z = relay.qnn.op.add(
+    z = relay.qnn.add(
         lhs=x,
         rhs=y,
         lhs_scale=relay.const(0.00784314, "float32"),

tests/python/relay/test_op_qnn_batch_matmul.py

@@ -182,7 +182,7 @@ def qnn_batch_matmul_driver(test_configuration):
     expected_out_dtype = test_configuration["out_dtype"]
     quantized_x = relay.var(quantized_x_name, shape=test_configuration["x_shape"], dtype=in_dtype)
     quantized_y = relay.var(quantized_y_name, shape=test_configuration["y_shape"], dtype=in_dtype)
-    mod = relay.qnn.op.batch_matmul(
+    mod = relay.qnn.batch_matmul(
         quantized_x,
         quantized_y,
         relay.const(test_configuration["x_zero_point"], "int32"),
@@ -192,7 +192,7 @@ def qnn_batch_matmul_driver(test_configuration):
     )
     if test_configuration["requantize"] is not None:
         requantize_config = test_configuration["requantize"]
-        mod = relay.qnn.op.requantize(
+        mod = relay.qnn.requantize(
             mod,
             input_scale=relay.const(requantize_config["input_scale"], "float32"),
             input_zero_point=relay.const(0, "int32"),

tests/python/relay/test_op_qnn_concatenate.py

@@ -34,7 +34,7 @@ def test_same_io_qnn_params():
 
     x = relay.var("x", shape=(1, 64), dtype=data_dtype)
     y = relay.var("y", shape=(1, 64), dtype=data_dtype)
-    z = relay.qnn.op.concatenate(
+    z = relay.qnn.concatenate(
         (x, y),
         input_scales=(x_scale, y_scale),
         input_zero_points=(zero, zero),
@@ -70,7 +70,7 @@ def test_different_io_qnn_params():
 
     x = relay.var("x", shape=(1, 64), dtype=data_dtype)
     y = relay.var("y", shape=(1, 64), dtype=data_dtype)
-    z = relay.qnn.op.concatenate(
+    z = relay.qnn.concatenate(
         (x, y),
         input_scales=(x_scale, y_scale),
         input_zero_points=(x_zero_point, y_zero_point),
@@ -106,7 +106,7 @@ def test_few_same_io_qnn_params():
 
     x = relay.var("x", shape=(1, 64), dtype=data_dtype)
     y = relay.var("y", shape=(1, 64), dtype=data_dtype)
-    z = relay.qnn.op.concatenate(
+    z = relay.qnn.concatenate(
         (x, y),
         input_scales=(x_scale, y_scale),
         input_zero_points=(x_zero_point, y_zero_point),
@@ -142,7 +142,7 @@ def test_same_i_qnn_params():
 
     x = relay.var("x", shape=(1, 64), dtype=data_dtype)
     y = relay.var("y", shape=(1, 64), dtype=data_dtype)
-    z = relay.qnn.op.concatenate(
+    z = relay.qnn.concatenate(
         (x, y),
         input_scales=(x_scale, y_scale),
         input_zero_points=(x_zero_point, y_zero_point),
@@ -177,7 +177,7 @@ def test_call_input():
     y_zero_point = relay.const(0, "int32")
 
     tup = relay.split(x, 2, axis=0)
-    z = relay.qnn.op.concatenate(
+    z = relay.qnn.concatenate(
         tup,
         input_scales=(x_scale, y_scale),
         input_zero_points=(x_zero_point, y_zero_point),