Get dim from the input of reshape and transpose (#3306)

Signed-off-by: Tung D. Le <[email protected]>

Author: tungld

src/Dialect/ONNX/ONNXOps/Canonicalize.cpp

@@ -1665,6 +1665,106 @@ struct RemoveDimZeroInputInConcatPattern
   }
 };
 
+// =============================================================================
+// Rewrite pattern onnx.dim
+// =============================================================================
+
+/// The pattern is to replace DimOp by a dimension value in the shape of
+/// ReshapeOp.
+///
+/// We would like to replace:
+/// ```
+/// %shape   = onnx.Concat(%d0, %d1, %d2)
+/// %reshape = onnx.Reshape(%X, %shape1) {allowzero = 0}
+/// %dim     = onnx.Dim(%reshape) {axis = 1}
+/// ```
+/// with
+/// ```
+/// %dim = %d1
+/// ```
+/// We only consider `allowzero=0` in this pattern.
+
+struct DimOpFromReshapeInputPattern : public OpRewritePattern<ONNXDimOp> {
+  using OpRewritePattern<ONNXDimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(
+      ONNXDimOp dimOp, PatternRewriter &rewriter) const final {
+    Value data = dimOp.getData();
+
+    // Donot handle unranked tensors.
+    if (!isRankedShapedType(data.getType()))
+      return rewriter.notifyMatchFailure(dimOp, "Input is unranked");
+
+    // Normalize axis.
+    int64_t rank = getRank(data.getType());
+    int64_t dimAxis = dimOp.getAxis();
+    if (dimAxis < 0)
+      dimAxis += rank;
+
+    // Dim is from a reshape op.
+    ONNXReshapeOp reshapeOp = data.getDefiningOp<ONNXReshapeOp>();
+    if (!reshapeOp)
+      return rewriter.notifyMatchFailure(dimOp, "Not found reshape op");
+    if (reshapeOp.getAllowzero() != 0)
+      return rewriter.notifyMatchFailure(dimOp, "Reshape op's allowzero != 0");
+
+    // Shape is from a concat op of dims.
+    Value shape = reshapeOp.getShape();
+    ONNXConcatOp concatOp = shape.getDefiningOp<ONNXConcatOp>();
+    if (!concatOp)
+      return rewriter.notifyMatchFailure(dimOp, "Not found concat op");
+    ValueRange dims = concatOp.getInputs();
+
+    // Ensure that the number of concat's inputs is equal to rank.
+    if (dims.size() != static_cast<uint64_t>(rank))
+      return rewriter.notifyMatchFailure(
+          dimOp, "Concat input size is not good");
+
+    // Values in shape can be -1 or 0 according to Reshape's definition.
+    // Those values are not real dimensions.
+    if (isConstOf(dims[dimAxis], -1) || (isConstOf(dims[dimAxis], 0)))
+      return rewriter.notifyMatchFailure(dimOp, "Dim at axis is -1 or 0");
+
+    rewriter.replaceOp(dimOp, dims[dimAxis]);
+    return success();
+  }
+};
+
+struct DimOpFromTransposeInputPattern : public OpRewritePattern<ONNXDimOp> {
+  using OpRewritePattern<ONNXDimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(
+      ONNXDimOp dimOp, PatternRewriter &rewriter) const final {
+    Value data = dimOp.getData();
+
+    // Donot handle unranked tensors.
+    if (!isRankedShapedType(data.getType()))
+      return rewriter.notifyMatchFailure(dimOp, "Input is unranked");
+
+    // Normalize axis.
+    int64_t rank = getRank(data.getType());
+    int64_t dimAxis = dimOp.getAxis();
+    if (dimAxis < 0)
+      dimAxis += rank;
+
+    // Dim is from a transpose op.
+    ONNXTransposeOp transposeOp = data.getDefiningOp<ONNXTransposeOp>();
+    if (!transposeOp)
+      return rewriter.notifyMatchFailure(dimOp, "Not found transpose op");
+    Value transposeData = transposeOp.getData();
+
+    // Transpose axes.
+    ArrayAttr permAttr = transposeOp.getPermAttr();
+    int64_t transposeAxis = ArrayAttrIntVal(permAttr, dimAxis);
+
+    Value replacedDim =
+        OnnxBuilder(rewriter, dimOp.getLoc()).dim(transposeData, transposeAxis);
+
+    rewriter.replaceOp(dimOp, replacedDim);
+    return success();
+  }
+};
+
 // =============================================================================
 // Rewrite pattern LayerNormalization
 // =============================================================================
@@ -2268,6 +2368,8 @@ void ONNXDropoutOp::getCanonicalizationPatterns(
 void ONNXDimOp::getCanonicalizationPatterns(
     RewritePatternSet &results, MLIRContext *context) {
   results.insert<DimOpToConstantPattern>(context);
+  results.insert<DimOpFromReshapeInputPattern>(context);
+  results.insert<DimOpFromTransposeInputPattern>(context);
 }
 
 /// on the ONNXEqualOp.

src/Dialect/ONNX/ONNXOps/OpHelper.cpp

@@ -597,6 +597,8 @@ WideNum asWideNum(double n, Type elemType) {
 /// Checks whether a constant tensor's elements are all equal to a given scalar.
 bool isConstOf(Value constValue, double n) {
   ElementsAttr constElements = getElementAttributeFromONNXValue(constValue);
+  if (!constElements)
+    return false;
   Type elemType = constElements.getElementType();
   assert(!elemType.isInteger(1) && "booleans are not supported");
   WideNum w = asWideNum(n, elemType);

test/mlir/onnx/onnx_canonicalization.mlir

@@ -1329,6 +1329,39 @@ func.func @test_dim_to_constant(%arg0: tensor<?x256xi64>) -> (tensor<1xi64>) {
 
 // -----
 
+func.func @dim_from_reshape(%arg0: tensor<1x32x?x?xf32>, %arg1: tensor<32x?x?xf32>) -> tensor<1xi64> {
+  %0 = onnx.Constant dense<32> : tensor<1xi64>
+  %1 = onnx.Constant dense<1> : tensor<1xi64>
+  %2 = "onnx.Dim"(%arg0) {axis = 2 : si64} : (tensor<1x32x?x?xf32>) -> tensor<1xi64>
+  %3 = "onnx.Dim"(%arg0) {axis = 3 : si64} : (tensor<1x32x?x?xf32>) -> tensor<1xi64>
+  %4 = "onnx.Concat"(%1, %0, %2, %3) {axis = 0 : si64} : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<4xi64>
+  %5 = "onnx.Reshape"(%arg1, %4) {allowzero = 0 : si64} : (tensor<32x?x?xf32>, tensor<4xi64>) -> tensor<1x32x?x?xf32>
+  %6 = "onnx.Dim"(%5) {axis = 2 : si64} : (tensor<1x32x?x?xf32>) -> tensor<1xi64>
+  return %6 : tensor<1xi64>
+
+// CHECK-LABEL:  func.func @dim_from_reshape
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<1x32x?x?xf32>, [[PARAM_1_:%.+]]: tensor<32x?x?xf32>) -> tensor<1xi64> {
+// CHECK:           [[VAR_0_:%.+]] = "onnx.Dim"([[PARAM_0_]]) {axis = 2 : si64} : (tensor<1x32x?x?xf32>) -> tensor<1xi64>
+// CHECK:           return [[VAR_0_]] : tensor<1xi64>
+// CHECK:         }
+}
+
+// -----
+
+func.func @dim_from_transpose(%arg0: tensor<1x32x?x?xf32>) -> tensor<1xi64> {
+  %0 = "onnx.Transpose"(%arg0) {perm = [0, 2, 1, 3]} : (tensor<1x32x?x?xf32>) -> tensor<1x?x32x?xf32>
+  %1 = "onnx.Dim"(%0) {axis = 1 : si64} : (tensor<1x?x32x?xf32>) -> tensor<1xi64>
+  return %1 : tensor<1xi64>
+
+// CHECK-LABEL:  func.func @dim_from_transpose
+// CHECK-SAME:   ([[PARAM_0_:%.+]]: tensor<1x32x?x?xf32>) -> tensor<1xi64> {
+// CHECK:           [[VAR_0_:%.+]] = "onnx.Dim"([[PARAM_0_]]) {axis = 2 : si64} : (tensor<1x32x?x?xf32>) -> tensor<1xi64>
+// CHECK:           return [[VAR_0_]] : tensor<1xi64>
+// CHECK:         }
+}
+
+// -----
+
 func.func @test_layout_transform(%arg0: tensor<5x3x32x32xf32, #onnx.layout<{dataLayout = "NCHW4C"}>>) -> tensor<5x3x32x32xf32, #onnx.layout<{dataLayout = "NCHW4C"}>> {
     %0 = "onnx.LayoutTransform"(%arg0) {target_layout = #onnx.layout<{dataLayout = "NCHW4C"}>} : (tensor<5x3x32x32xf32,#onnx.layout<{dataLayout = "NCHW4C"}>>) -> tensor<5x3x32x32xf32, #onnx.layout<{dataLayout = "NCHW4C"}>>
     onnx.Return %0 : tensor<5x3x32x32xf32, #onnx.layout<{dataLayout = "NCHW4C"}>>