Merge pull request #442 from Xilinx/planzase.AIESW-13111.gather_check_for_static_shape

[AIESW-13111] Check Gather input shape is static

Author: jorickert

src/Conversion/ONNXToTOSA/Tensor/Gather.cpp

@@ -17,6 +17,7 @@
 #include "src/Conversion/ONNXToTOSA/ONNXToTOSACommon.hpp"
 #include "src/Conversion/ONNXToTOSA/ONNXToTOSALegalizeUtils.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
+#include "src/Support/TypeUtilities.hpp"
 #include "llvm/ADT/SmallVector.h"
 
 using namespace mlir;
@@ -46,7 +47,8 @@ class ONNXGatherLoweringToTOSA : public OpConversionPattern<ONNXGatherOp> {
     if (!onnx_mlir::isRankedShapedType(inputType))
       return rewriter.notifyMatchFailure(op, "input is not a ranked tensor");
 
-    if (!hasStaticShape(result.getType()))
+    if (!hasStaticShape(inputType) || !hasStaticShape(indices.getType()) ||
+        !hasStaticShape(result.getType()))
       return rewriter.notifyMatchFailure(op, "dynamic shapes not supported");
 
     auto resultTy = dyn_cast<TensorType>(op.getType());

src/Dialect/ONNX/Transforms/SimplifyShapeRelatedOps.cpp

@@ -65,7 +65,6 @@ Now, it's straightforward to update the output shape of Reshape from
 #include "src/Dialect/ONNX/ONNXOps/OpHelper.hpp"
 #include "src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp"
 #include "src/Pass/Passes.hpp"
-#include "src/Support/TypeUtilities.hpp"
 
 #define DEBUG_TYPE "simplify_shape_related_ops"
 

test/mlir/conversion/onnx_to_tosa/Tensor/Gather.mlir

@@ -175,3 +175,21 @@ func.func @test_gather_dynamic_shape_indices_i32(%arg0 : tensor<?x4xf32>, %indic
 // CHECK-LABEL: test_gather_dynamic_shape_indices_i32
 // CHECK: onnx.Gather
 }
+
+// -----
+
+func.func @test_gather_dynamic_input_static_output(%arg0 : tensor<?x2xf32>, %indices: tensor<?xi64>) -> tensor<1x2xf32> {
+  %0 = "onnx.Gather"(%arg0, %indices) {axis = 0 : si64, onnx_node_name = "/Gather_16"} : (tensor<?x2xf32>, tensor<?xi64>) -> tensor<1x2xf32>
+  "func.return"(%0) : (tensor<1x2xf32>) -> ()
+// CHECK-LABEL: test_gather_dynamic_input_static_output
+// CHECK: onnx.Gather
+}
+
+// -----
+
+func.func @test_gather_dynamic_indices(%arg0 : tensor<1x2xf32>, %indices: tensor<?xi64>) -> tensor<1x2xf32> {
+  %0 = "onnx.Gather"(%arg0, %indices) {axis = 0 : si64, onnx_node_name = "/Gather_16"} : (tensor<1x2xf32>, tensor<?xi64>) -> tensor<1x2xf32>
+  "func.return"(%0) : (tensor<1x2xf32>) -> ()
+// CHECK-LABEL: test_gather_dynamic_indices
+// CHECK: onnx.Gather
+}