Dynamic ONNX Importer (apache#6351)

* Change onnx importer to use dynamic upsampling3d (neo-ai#3)

fix pylint

* Refactor ONNX frontend to be dynamic

Make OneHot dynamic

Support BatchMatMul with dynamically shaped inputs

fix dynamic broadcast

Add null checks to broadcast_to rel functions

fail more isolated broadcast_to test

use StructuralEqual instead of pointer comparisions in dynamic_to_static pass

add an optional weight freeze argument to onnx importer

convert onnx resize to dynamic op

add dynamic expand to onnx importer

add a shape_func for power

fix BERTSquad, lint

handle onnx graph initializer parameters more intelligently

* Dynamic ONNX importer: Upsampling and Pad (neo-ai#2)

fix lint

fix Call reference

fix a type issue with expand

fix a bad test refactor

respond to review comments, fix batch matmul tests

* black format

* fix batch matmul test

* add dynamic strided slice to the onnx importer

* fix clip importer

* fix qnn tutorial

* fix bad merge, respond to review comments

* add a simple dynamic model test

* Add dynamic-shaped autopadding to convolution and pooling ops

* fix dynamic issues in a few ops

* fix pylint

* disable tests onnxrt doesn't support

* fix pytorch test

* respond to review comments

* add documentation about partially supporting dynamic shapes

Co-authored-by: Lily Orth-Smith <[email protected]>

python/tvm/relay/frontend/onnx.py

@@ -528,6 +528,15 @@ def flatten_to_3d(x, x_shape):
             # Convert a and b into 3 dimensional tensors.
             a = flatten_to_3d(inputs[0], a_shape)
             b = flatten_to_3d(inputs[1], b_shape)
+            # Broadcast b to match batch size of a
+            new_b_shape = _op.concatenate(
+                [
+                    _op.strided_slice(_op.shape_of(a), [0], [1]),
+                    _op.strided_slice(_op.shape_of(b), [1], [3]),
+                ],
+                0,
+            )
+            b = _op.broadcast_to(b, new_b_shape)
             # Transpose matrix dimensions of b.
             b = _op.transpose(b, [0, 2, 1])
             # Perform a batch matmul.
@@ -1893,19 +1902,6 @@ def _impl_v1(cls, inputs, attr, params):
         return _op.topk(inputs[0], inputs[1], axis=axis)
 
 
-class Range(OnnxOpConverter):
-    """Operator converter for Range"""
-
-    @classmethod
-    def _impl_v1(cls, inputs, attr, params):
-        if len(inputs) != 3:
-            raise ValueError("Expect 3 input only")
-
-        return _op.arange(
-            inputs[0], inputs[1], inputs[2], dtype=infer_type(inputs[0]).checked_type.dtype
-        )
-
-
 class MaxRoiPool(OnnxOpConverter):
     """Operator converter for MaxRoiPool."""
 

python/tvm/topi/nn/batch_matmul.py

@@ -41,21 +41,19 @@ def batch_matmul(x, y, oshape=None):
     output : tvm.te.Tensor
         3-D with shape [batch, M, N]
     """
-    assert len(x.shape) == 3 and len(y.shape) == 3, "only support 3-dim batch_matmul"
-    x_shape = get_const_tuple(x.shape)
-    y_shape = get_const_tuple(y.shape)
-    XB = x_shape[0]
-    YB = y_shape[0]
-    _, M, K = x.shape
-    k = te.reduce_axis((0, K), name="k")
     if oshape is None:
-        assert XB == YB or XB == 1 or YB == 1, "batch dimension doesn't match"
+        assert len(x.shape) == 3 and len(y.shape) == 3, "only support 3-dim batch_matmul"
+        x_shape = get_const_tuple(x.shape)
+        y_shape = get_const_tuple(y.shape)
+        assert x_shape[0] == y_shape[0], "batch dimension doesn't match"
         assert x_shape[2] == y_shape[2], "shapes of x and y is inconsistant"
-        batch = max(XB, YB)
+        batch, M, K = x.shape
         N = y.shape[1]
+        k = te.reduce_axis((0, K), name="k")
         oshape = (batch, M, N)
+    else:
+        _, _, K = x.shape
+        k = te.reduce_axis((0, K), name="k")
     return te.compute(
-        oshape,
-        lambda b, i, j: te.sum(x[b if XB != 1 else 0, i, k] * y[b if YB != 1 else 0, j, k], axis=k),
-        tag="batch_matmul",
+        oshape, lambda b, i, j: te.sum(x[b, i, k] * y[b, j, k], axis=k), tag="batch_matmul"
     )

src/relay/op/nn/convolution.h

@@ -258,6 +258,9 @@ bool Conv2DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
         return false;
       }
     }
+    if (!dshape_nchw[1].as<tir::AnyNode>() && !wshape[1].as<tir::AnyNode>()) {
+      CHECK(reporter->AssertEQ(indexdiv(dshape_nchw[1], param->groups), wshape[1]));
+    }
     channels = wshape[0];
     dilated_ksize_y = 1 + (wshape[2] - 1) * param->dilation[0];
     dilated_ksize_x = 1 + (wshape[3] - 1) * param->dilation[1];
@@ -369,7 +372,6 @@ bool Conv3DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
           << "Conv3D: shape of weight is inconsistent with channels, "
           << " channels=" << param->channels << " wshape=" << wshape;
     }
-
     if (!dshape_ncdhw[1].as<tir::AnyNode>() && !wshape[1].as<tir::AnyNode>()) {
       CHECK(reporter->AssertEQ(indexdiv(dshape_ncdhw[1], param->groups), wshape[1]));
     }

src/relay/op/nn/nn.cc

@@ -863,9 +863,8 @@ bool BatchMatmulRel(const Array<Type>& types, int num_inputs, const Attrs& attrs
     }
   }
   if (!is_dyn) {
-    CHECK(reporter->AssertEQ(x->shape[0], y->shape[0]) || reporter->AssertEQ(x->shape[0], 1) ||
-          reporter->AssertEQ(y->shape[0], 1))
-        << "BatchDot: batch dimensions don't match, "
+    CHECK(reporter->AssertEQ(x->shape[0], y->shape[0]))
+        << "BatchDot: batch dimension doesn't match, "
         << " x shape=" << x->shape << ", y shape=" << y->shape;
     CHECK(reporter->AssertEQ(x->shape[2], y->shape[2]))
         << "BatchDot: shapes of x and y is inconsistent, "