Split inferShapes and added verify to Flatten operator (llvm#1140)

* Split inferShapes and added verifier to Flatten operators
Signed-off-by: Adrian Sion <[email protected]>

src/Dialect/ONNX/ONNXOps.cpp

@@ -2835,19 +2835,37 @@ LogicalResult ONNXSplitV11Op::inferShapes(
 // Flatten
 //===----------------------------------------------------------------------===//
 
+static LogicalResult verify(ONNXFlattenOp op) {
+
+  if (!hasShapeAndRank(op.input())) {
+    return success();
+  }
+  auto inTy = op.input().getType().dyn_cast<ShapedType>();
+  if (!inTy) {
+    return success();
+  }
+
+  int64_t axisValue = op.axis();
+  auto inputShape = inTy.getShape();
+  int64_t inputRank = inputShape.size();
+
+  if (axisValue < -1 * inputRank || axisValue > inputRank) {
+    return op.emitError("ONNXFlattenOP: axis() value is out of range");
+  }
+
+  return success();
+}
+
 LogicalResult ONNXFlattenOp::inferShapes(
     std::function<void(mlir::Region &)> doShapeInference) {
-  auto inTy = input().getType().dyn_cast<RankedTensorType>();
+  auto inTy = input().getType().dyn_cast_or_null<RankedTensorType>();
   if (!inTy) {
     return success();
   }
 
   int64_t axisValue = axis();
   auto inputShape = inTy.getShape();
   int64_t inputRank = inputShape.size();
-  if (axisValue < -1 * inputRank || axisValue > inputRank) {
-    return emitOpError("ONNXFlattenOP: axis() value is out of range");
-  }
 
   SmallVector<int64_t, 2> dims;
 

src/Dialect/ONNX/ONNXOps.td.inc

@@ -1502,6 +1502,7 @@ def ONNXFlattenOp:ONNX_Op<"Flatten",
       return {20};
     }
   }];
+  let verifier = [{ return ::verify(*this); }];
 }
 
 def ONNXFloorOp:ONNX_Op<"Floor",

test/mlir/onnx/invalid.mlir

@@ -40,3 +40,13 @@ func @test_concat_verifier_3(%arg0 : tensor<5x5x1x32xf32>, %arg1 : tensor<5x5x3x
   %1 = "onnx.Concat"(%arg0, %arg1, %arg2) { axis = 1 : si64} : (tensor<5x5x1x32xf32>, tensor<5x5x3x32xf32>, tensor<5x5x5x32xf32>)  -> tensor<*xf32>
   "std.return"(%1) : (tensor<*xf32>) -> ()
 }
+
+// -----
+
+func @test_flatten_verifier_1(%arg0 : tensor<5x5x1x32xf32>) -> tensor<*xf32> {
+  // expected-error @+1 {{ONNXFlattenOP: axis() value is out of range}}
+  %1 = "onnx.Flatten"(%arg0) { axis = 5 : si64} : (tensor<5x5x1x32xf32>) -> tensor<*xf32>
+  "std.return"(%1) : (tensor<*xf32>) -> ()
+}
+
+// -----

utils/gen_onnx_mlir.py

@@ -301,6 +301,7 @@
     'Conv',
     'DepthToSpace',
     'Expand',
+    'Flatten',
     'Hardmax',
     'InstanceNormalization',
     'Mod',