[Mhlo] fixes for code review

Signed-off-by: chongsong.chen <chongsong.chen@bytedance.com>
onnx · Jul 15, 2022 · e3ed189 · e3ed189
1 parent 2cedfc1
commit e3ed189
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 15 deletions.
diff --git a/src/Conversion/ONNXToMhlo/Tensor/Concat.cpp b/src/Conversion/ONNXToMhlo/Tensor/Concat.cpp
@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "src/Conversion/ONNXToMhlo/ONNXToMhloCommon.hpp"
+#include "src/Support/TypeUtilities.hpp"
 
 using namespace mlir;
 
@@ -31,17 +32,13 @@ struct ONNXConcatOpLoweringToMhlo : public ConversionPattern {
     ONNXConcatOpAdaptor operandAdaptor(operands);
     ONNXConcatOp concatOp = llvm::cast<ONNXConcatOp>(op);
 
-    if (op->getNumResults() < 1) {
-      op->emitError() << "ONNXConcatOp Has No Output\n";
-      return failure();
-    }
-    RankedTensorType resultType =
-        op->getResult(0).getType().dyn_cast_or_null<RankedTensorType>();
-    if (resultType == nullptr) {
+    assert(op->getNumResults() == 1 && "ONNXConcatOp shoule have 1 result");
+    Type resultType = op->getResult(0).getType();
+    if (!onnx_mlir::isRankedShapedType(resultType)) {
       op->emitError() << "Concat Output Is Not Ranked\n";
       return failure();
     }
-    int64_t rank = resultType.getRank();
+    int64_t rank = onnx_mlir::getRank(resultType);
     int64_t axis = concatOp.axis();
     axis = axis >= 0 ? axis : rank + axis;
     assert(axis >= -rank && axis <= rank - 1 && "Axis out of rank range");

diff --git a/src/Support/TypeUtilities.cpp b/src/Support/TypeUtilities.cpp
@@ -36,6 +36,12 @@ ArrayRef<int64_t> getShape(Type ty) {
   return ty.cast<ShapedType>().getShape();
 }
 
+/// Get rank.
+int64_t getRank(Type ty) {
+  assert(isRankedShapedType(ty) && "Type must be ranked");
+  return ty.cast<ShapedType>().getRank();
+}
+
 /// Get the number of elements.
 int64_t getNumberOfElements(Type ty) {
   ArrayRef<int64_t> shape = getShape(ty);

diff --git a/src/Support/TypeUtilities.hpp b/src/Support/TypeUtilities.hpp
@@ -22,6 +22,8 @@ mlir::Type getElementType(mlir::Type ty);
 bool isRankedShapedType(mlir::Type ty);
 /// Get shape.
 llvm::ArrayRef<int64_t> getShape(mlir::Type ty);
+/// Get rank.
+int64_t getRank(mlir::Type ty);
 /// Get the number of elements.
 int64_t getNumberOfElements(mlir::Type ty);
 /// Get the element size in bytes.

diff --git a/test/mlir/conversion/onnx_to_mhlo/Tensor/Concat.mlir b/test/mlir/conversion/onnx_to_mhlo/Tensor/Concat.mlir
@@ -1,13 +1,13 @@
 // RUN: onnx-mlir-opt --convert-onnx-to-mhlo %s -split-input-file | FileCheck %s
 
 // Test when output shape is unkown
-func @test_concat_unknown_dims(%arg0 : tensor<5x5x1x32xf32>, %arg1 : tensor<5x5x3x32xf32>) -> tensor<?x?x?x?xf32> {
-  %0 = "onnx.Concat"(%arg0, %arg1) { axis = 2 : si64} : (tensor<5x5x1x32xf32>, tensor<5x5x3x32xf32>) -> tensor<?x?x?x?xf32>
-  "func.return"(%0) : (tensor<?x?x?x?xf32>) -> ()
-// CHECK-LABEL:  func @test_concat_unknown_dims
-// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<5x5x1x32xf32>, [[PARAM_1_:%.+]]: tensor<5x5x3x32xf32>) -> tensor<?x?x?x?xf32> {
-// CHECK-NEXT:    [[VAR_0_:%.+]] = "mhlo.concatenate"(%arg0, %arg1) {dimension = 2 : i64} : (tensor<5x5x1x32xf32>, tensor<5x5x3x32xf32>) -> tensor<?x?x?x?xf32>
-// CHECK-NEXT:   return [[VAR_0_]] : tensor<?x?x?x?xf32>
+func @test_concat_dynamic_shape(%arg0 : tensor<5x5x?x32xf32>, %arg1 : tensor<5x5x?x32xf32>) -> tensor<5x5x?x32xf32> {
+  %0 = "onnx.Concat"(%arg0, %arg1) { axis = 2 : si64} : (tensor<5x5x?x32xf32>, tensor<5x5x?x32xf32>) -> tensor<5x5x?x32xf32>
+  "func.return"(%0) : (tensor<5x5x?x32xf32>) -> ()
+// CHECK-LABEL:  func @test_concat_dynamic_shape
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<5x5x?x32xf32>, [[PARAM_1_:%.+]]: tensor<5x5x?x32xf32>) -> tensor<5x5x?x32xf32> {
+// CHECK-NEXT:    [[VAR_0_:%.+]] = "mhlo.concatenate"(%arg0, %arg1) {dimension = 2 : i64} : (tensor<5x5x?x32xf32>, tensor<5x5x?x32xf32>) -> tensor<5x5x?x32xf32>
+// CHECK-NEXT:   return [[VAR_0_]] : tensor<5x5x?x32xf32>
 // CHECK-NEXT:   }
 }