int32 pass

Signed-off-by: Chen Tong <[email protected]>

src/Conversion/ONNXToKrnl/Tensor/GatherND.cpp

@@ -74,6 +74,8 @@ struct ONNXGatherNDOpLowering : public OpConversionPattern<ONNXGatherNDOp> {
     auto indicesType = indices.getType().cast<ShapedType>();
     ArrayRef<int64_t> indicesShape = indicesType.getShape();
     int64_t indicesLastDim = indicesShape[indicesRank - 1];
+    // ToFix: Handle case in which indicesLastDim is kDynamic.
+    // Currently, such case is detected by ONNXPreKrnlVerifyPass.
     assert((indicesLastDim >= 1 && indicesLastDim <= dataRank - b) &&
            "indices.shape[-1] must be in the range [1, dataRank - b]");
 

src/Dialect/ONNX/ONNXOps/Tensor/GatherND.cpp

@@ -50,6 +50,8 @@ LogicalResult ONNXGatherNDOpShapeHelper::computeShape() {
   assert(b >= 0 && "batch_dim should not be negative");
   assert(b < std::min(dataRank, indicesRank) &&
          "batch_dims must be smaller than the min(dataRank, indicesRank)");
+  // ToFix: Handle the case when indicesLastDim is dynamic and the rank
+  // of output tensor is known.
   assert((indicesLastDim >= 1 && indicesLastDim <= dataRank - b) &&
          "indices.shape[-1] must be in the range [1, dataRank - b]");
 

src/Dialect/ONNX/Transforms/ONNXPreKrnlVerifyPass.cpp

@@ -25,6 +25,7 @@
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 #include "src/Interface/ShapeInferenceOpInterface.hpp"
 #include "src/Pass/Passes.hpp"
+#include "src/Support/TypeUtilities.hpp"
 
 using namespace mlir;
 
@@ -70,6 +71,15 @@ class ONNXPreKrnlVerifyPass : public mlir::PassWrapper<ONNXPreKrnlVerifyPass,
       } else if (!ty.isa<RankedTensorType>() && !ty.isa<NoneType>()) {
         op.emitError("not ranked");
         return failure();
+      } else if (ONNXGatherNDOp gatherNDOp = llvm::dyn_cast<ONNXGatherNDOp>(op)) {
+	Value indices = gatherNDOp.getIndices();
+	Type indicesType = indices.getType();
+	ArrayRef<int64_t> indicesShape = onnx_mlir::getShape(indices.getType());
+	int64_t indicesRank = onnx_mlir::getRank(indicesType);
+	if (indicesShape[indicesRank - 1] == ShapedType::kDynamic) {
+	  op.emitError("last dim of indices in GatherND is dynamic");
+	  return failure();
+	}
       }
     }
     return success();

test/backend/inference_backend.py

@@ -1092,22 +1092,27 @@ def get_test_models():
         },
         # ==OP== GatherND
         # ==MIN== 11
+        # According to onnx specification:
+        # output tensor of rank q + r - indices_shape[-1] - 1 - b
+        # onnx-mlir can not infer the shape if the last dim of indices
+        # input is dynamic. Therefore, {0: {-1}} is used for DYNAMIC_SHAPE
         "test_gathernd_example_int32_cpu": {
             STATIC_SHAPE: {},
-            # Issue #2639: Dynamic test fails. Need to be fixed.
-            # DYNAMIC_SHAPE: {-1: {-1}},
+            DYNAMIC_SHAPE: {0: {-1}},
             CONSTANT_INPUT: {-1},
         },
         "test_gathernd_example_float32_cpu": {
             STATIC_SHAPE: {},
             # Issue #2639: Dynamic test fails. Need to be fixed.
-            # DYNAMIC_SHAPE: {-1: {-1}},
+            # Dynamic test failed due to result difference
+            # DYNAMIC_SHAPE: {0: {-1}},
             CONSTANT_INPUT: {-1},
         },
         "test_gathernd_example_int32_batch_dim1_cpu": {
             STATIC_SHAPE: {},
             # Issue #2639: Dynamic test fails. Need to be fixed.
-            # DYNAMIC_SHAPE: {-1: {-1}},
+            # Dynamic test failed due to result difference
+            # DYNAMIC_SHAPE: {0: {-1}},
             CONSTANT_INPUT: {-1},
         },
         # ==OP== Gelu