[MLIR][Arith] Fold trunci with ext if the bit width of the input type of ext is greater than the bit width of the result type of trunci

[Miss_Grape]

Summary

This patch is mainly to deal with folding trunci with ext，as flows：

trunci(zexti(a)) -> trunci(a)
trunci(sexti(a)) -> trunci(a)

[ftynse]

Please provide a commit description that explains what you do and, most importantly, why -- https://mlir.llvm.org/getting_started/Contributing/#commit-messages.

Also naming nit: the code you are editing is a folding, not a canonicalization.

[ftynse]

mlir/test/Dialect/Arith/canonicalize.mlir

432

Nit: looking at the name, I'd expect to see trunc(ext(trunc)) in the body.

[Miss_Grape]

address the comment

[ftynse]

The commit description still doesn't describe the details of what it is doing and why, only confuses the reader. The code before this commit *was folding trunc(ext)"! This commit does *not* add that. It does, however, *fix* a case that was mishandled. Please find the appropriate way to express that.

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

1245–1259

The following

Type dstType = getElementTypeOrSelf(getType());
Value src = getOperand().getDefiningOp()->getOperand(0);
Type srcType = src.getType();
if (srcType.cast<IntegerType>().getWIdth() >
    dstType.cast<IntegerType>().getWidth()) {
  setOperand(src);
  return getResult();
}
return src;

looks more readable, concise and avoids doing the same thing twice using different APIs (querying the kind of the defining op, getting its operand).

[Miss_Grape]

address the comment

[ftynse]

The actual comment was not addressed:

The commit description still doesn't describe the details of what it is doing and why, only confuses the reader. The code before this commit *was folding trunc(ext)"! This commit does *not* add that. It does, however, *fix* a case that was mishandled. Please find the appropriate way to express that.

[Miss_Grape]

In D140604#4039078, @ftynse wrote:

The actual comment was not addressed:

The commit description still doesn't describe the details of what it is doing and why, only confuses the reader. The code before this commit *was folding trunc(ext)"! This commit does *not* add that. It does, however, *fix* a case that was mishandled. Please find the appropriate way to express that.

address the comments

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

 
 //===----------------------------------------------------------------------===//
 // TruncIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::TruncIOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 1 && "unary operation takes one operand");
 
+  if (matchPattern(getOperand(), m_Op<arith::ExtUIOp>()) ||
+      matchPattern(getOperand(), m_Op<arith::ExtSIOp>())) {
+    Value src = getOperand().getDefiningOp()->getOperand(0);
+    Type srcType = getElementTypeOrSelf(src.getType());
+    Type dstType = getElementTypeOrSelf(getType());
+    // trunci(zexti(a)) -> trunci(a)
+    // trunci(sexti(a)) -> trunci(a)
+    if (srcType.cast<IntegerType>().getWidth() >
+        dstType.cast<IntegerType>().getWidth()) {
+      setOperand(src);
+      return getResult();
+    }
     // trunci(zexti(a)) -> a
     // trunci(sexti(a)) -> a
-  if (matchPattern(getOperand(), m_Op<arith::ExtUIOp>()) ||
-      matchPattern(getOperand(), m_Op<arith::ExtSIOp>()))
-    return getOperand().getDefiningOp()->getOperand(0);
+    return src;
+  }
 
   // trunci(trunci(a)) -> trunci(a))
   if (matchPattern(getOperand(), m_Op<arith::TruncIOp>())) {
     setOperand(getOperand().getDefiningOp()->getOperand(0));
     return getResult();
   }
 
   Type resType = getElementTypeOrSelf(getType());

mlir/test/Dialect/Arith/canonicalize.mlir

 //       CHECK-NOT:  trunci
 //       CHECK:   return  %arg0
 func.func @truncExtui(%arg0: i32) -> i32 {
   %extui = arith.extui %arg0 : i32 to i64
   %trunci = arith.trunci %extui : i64 to i32
   return %trunci : i32
 }
 
+// CHECK-LABEL: @truncExtui2
+//       CHECK:  %[[ARG0:.+]]: i32
+//       CHECK:  %[[CST:.*]] = arith.trunci %[[ARG0:.+]] : i32 to i16
+//       CHECK:   return  %[[CST:.*]]
+func.func @truncExtui2(%arg0: i32) -> i16 {
+  %extui = arith.extui %arg0 : i32 to i64
+  %trunci = arith.trunci %extui : i64 to i16
+  return %trunci : i16
+}
+
+// CHECK-LABEL: @truncExtuiVector
+//       CHECK:  %[[ARG0:.+]]: vector<2xi32>
+//       CHECK:  %[[CST:.*]] = arith.trunci %[[ARG0:.+]] : vector<2xi32> to vector<2xi16>
+//       CHECK:   return  %[[CST:.*]]
+func.func @truncExtuiVector(%arg0: vector<2xi32>) -> vector<2xi16> {
+  %extsi = arith.extui %arg0 : vector<2xi32> to vector<2xi64>
+  %trunci = arith.trunci %extsi : vector<2xi64> to vector<2xi16>
+  return %trunci : vector<2xi16>
+}
+
 // CHECK-LABEL: @truncExtsi
 //       CHECK-NOT:  trunci
 //       CHECK:   return  %arg0
 func.func @truncExtsi(%arg0: i32) -> i32 {
   %extsi = arith.extsi %arg0 : i32 to i64
   %trunci = arith.trunci %extsi : i64 to i32
   return %trunci : i32
 }
 
+// CHECK-LABEL: @truncExtsi2
+//       CHECK:  %[[ARG0:.+]]: i32
+//       CHECK:  %[[CST:.*]] = arith.trunci %[[ARG0:.+]] : i32 to i16
+//       CHECK:   return  %[[CST:.*]]
+func.func @truncExtsi2(%arg0: i32) -> i16 {
+  %extsi = arith.extsi %arg0 : i32 to i64
+  %trunci = arith.trunci %extsi : i64 to i16
+  return %trunci : i16
+}
+
+// CHECK-LABEL: @truncExtsiVector
+//       CHECK:  %[[ARG0:.+]]: vector<2xi32>
+//       CHECK:  %[[CST:.*]] = arith.trunci %[[ARG0:.+]] : vector<2xi32> to vector<2xi16>
+//       CHECK:   return  %[[CST:.*]]
+func.func @truncExtsiVector(%arg0: vector<2xi32>) -> vector<2xi16> {
+  %extsi = arith.extsi %arg0 : vector<2xi32> to vector<2xi64>
+  %trunci = arith.trunci %extsi : vector<2xi64> to vector<2xi16>
+  return %trunci : vector<2xi16>
+}
+
 // CHECK-LABEL: @truncConstantSplat
 //       CHECK:   %[[cres:.+]] = arith.constant dense<-2> : vector<4xi8>
 //       CHECK:   return %[[cres]]
 func.func @truncConstantSplat() -> vector<4xi8> {
   %c-2 = arith.constant -2 : i16
   %splat = vector.splat %c-2 : vector<4xi16>
   %trunc = arith.trunci %splat : vector<4xi16> to vector<4xi8>
   return %trunc : vector<4xi8>