[MLIR][STD] Fold trunci (zexti).

[ergawy]

Summary

This patch folds the following pattern:

%arg0 = ...
%0 = zexti %arg0 : i1 to i8
%1 = trunci %0 : i8 to i1

into just %arg0.

[ergawy]

I came across this pattern while working on detensoring in linalg-on-tensors for IREE. Let me know if there are any reasons why this might be a bad idea.

[mehdi_amini]

Should be valid with sexti as well right?

If the bit-widths of trunci's destination and source types do not match the bit-widths of zexti's source and destination types, respectively, no folding takes place.

Your folder does not do that actually. It is just the framework that does not perform such replacement for all folders.

(side note, you could write simply : If the type of trunci's destination does not match zexti's source types, no folding takes place. ; the intermediate type does not matter here)

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2185

Can you add a one line comment: trunci(zexti(a)) -> a

[mehdi_amini]

mlir/test/Transforms/canonicalize.mlir

1095

This could be canonicalized to zexti %arg0 : i1 to i2 directly

[ergawy]

• Handle review comments.
• Rebase.

[ergawy]

Your folder does not do that actually. It is just the framework that does not perform such replacement for all folders.
(side note, you could write simply : If the type of trunci's destination does not match zexti's source types, no folding takes place. ; the intermediate type does not matter here)

I removed this part of the commit message to describe the changes introduced more accurately.

mlir/test/Transforms/canonicalize.mlir

1095

Added a TODO. I will address this it in a follow up partch.

[ergawy]

Should be valid with sexti as well right?

Will also do in a follow up patch.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let builders = [
     }]>];
 
   let parser = [{
     return impl::parseCastOp(parser, result);
   }];
   let printer = [{
     return printStandardCastOp(this->getOperation(), p);
   }];
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // UIToFPOp
 //===----------------------------------------------------------------------===//
 
 def UIToFPOp : ArithmeticCastOp<"uitofp">, Arguments<(ins AnyType:$in)> {
   let summary = "cast from unsigned integer type to floating-point";

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

static LogicalResult verify(TruncateIOp op) {
   if (srcType.cast<IntegerType>().getWidth() <=
       dstType.cast<IntegerType>().getWidth())
     return op.emitError("operand type ")
            << srcType << " must be wider than result type " << dstType;
 
   return success();
 }
 
+OpFoldResult TruncateIOp::fold(ArrayRef<Attribute> operands) {
+  // trunci(zexti(a)) -> a
+  if (matchPattern(getOperand(), m_Op<ZeroExtendIOp>()))
+    return getOperand().getDefiningOp()->getOperand(0);
+
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // UnsignedDivIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult UnsignedDivIOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 2 && "binary operation takes two operands");
 
   // Don't fold if it would require a division by zero.

mlir/test/Transforms/canonicalize.mlir

func @subtensor(%t: tensor<8x16x4xf32>, %arg0 : index, %arg1 : index)
   // CHECK-SAME: tensor<7x11x2xf32> to tensor<2x?x2xf32>
   // CHECK: tensor.cast %{{.*}} : tensor<2x?x2xf32> to tensor<?x?x?xf32>
   %2 = subtensor %1[%c0, %c0, %c0] [%c2, %arg0, %c2] [%c1, %c1, %c1]
     : tensor<?x?x?xf32> to tensor<?x?x?xf32>
 
   return %2 : tensor<?x?x?xf32>
 }
 
+// -----
+
+// CHECK-LABEL: func @fold_trunci
+// CHECK-SAME:    (%[[ARG0:[0-9a-z]*]]: i1)
+func @fold_trunci(%arg0: i1) -> i1 attributes {} {
+  // CHECK-NEXT: return %[[ARG0]] : i1
+  %0 = zexti %arg0 : i1 to i8
+  %1 = trunci %0 : i8 to i1
+  return %1 : i1
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_trunci_vector
+// CHECK-SAME:    (%[[ARG0:[0-9a-z]*]]: vector<4xi1>)
+func @fold_trunci_vector(%arg0: vector<4xi1>) -> vector<4xi1> attributes {} {
+  // CHECK-NEXT: return %[[ARG0]] : vector<4xi1>
+  %0 = zexti %arg0 : vector<4xi1> to vector<4xi8>
+  %1 = trunci %0 : vector<4xi8> to vector<4xi1>
+  return %1 : vector<4xi1>
+}
+
+// -----
+
+// TODO Canonicalize this into:
+//   zexti %arg0 : i1 to i2
+
+// CHECK-LABEL: func @do_not_fold_trunci
+// CHECK-SAME:    (%[[ARG0:[0-9a-z]*]]: i1)
+func @do_not_fold_trunci(%arg0: i1) -> i2 attributes {} {
+  // CHECK-NEXT: zexti %[[ARG0]] : i1 to i8
+  // CHECK-NEXT: %[[RES:[0-9a-z]*]] = trunci %{{.*}} : i8 to i2
+  // CHECK-NEXT: return %[[RES]] : i2
+  %0 = zexti %arg0 : i1 to i8
+  %1 = trunci %0 : i8 to i2
+  return %1 : i2
+}
+
+// -----
+
+// CHECK-LABEL: func @do_not_fold_trunci_vector
+// CHECK-SAME:    (%[[ARG0:[0-9a-z]*]]: vector<4xi1>)
+func @do_not_fold_trunci_vector(%arg0: vector<4xi1>) -> vector<4xi2> attributes {} {
+  // CHECK-NEXT: zexti %[[ARG0]] : vector<4xi1> to vector<4xi8>
+  // CHECK-NEXT: %[[RES:[0-9a-z]*]] = trunci %{{.*}} : vector<4xi8> to vector<4xi2>
+  // CHECK-NEXT: return %[[RES]] : vector<4xi2>
+  %0 = zexti %arg0 : vector<4xi1> to vector<4xi8>
+  %1 = trunci %0 : vector<4xi8> to vector<4xi2>
+  return %1 : vector<4xi2>
+}