[MLIR] Canonicalization of Integer Cast Operations

[wsmoses]

Summary

Canonicalize/fold SExt and IndexCast

[bondhugula]

Please expand the commit summary to a line or two at least. It really doesn't summarize the canonicalization part.

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

1245

Missing doc comment.

1471–1472

Please use dyn_cast - will lead to one check less with assertions enabled and more compact.

2731–2732

Likewise.

mlir/test/Dialect/Standard/canonicalize.mlir

406–410

This looks wrong to me. You don't get the same bits in index without doing the sign extension?

431

Missing new line at end of file.

[wsmoses]

Use dyn_cast and add comments to commit

[wsmoses]

The commit message has been revised to the following (though apparently not updating on Phabricator summary):

Author: William S. Moses <[e-mail]>
Date:   Sun May 2 00:39:45 2021 -0400

    [MLIR] Canonicalization of Integer Cast Operations
    
    1) Canonicalize IndexCast(SExt(x)) => IndexCast(x)
    2) Provide constant folds of sign_extend and truncate
    
    Differential Revision: https://reviews.llvm.org/D101714

mlir/test/Dialect/Standard/canonicalize.mlir

406–410

From my reading of the lang ref (https://mlir.llvm.org/docs/Dialects/Standard/#stdindex_cast-indexcastop), this should be correct. Namely, index_cast is itself defined to be a sign extension so it is fine to fold the other extension into it.

[bondhugula]

The phab summary will have to be updated manually.

mlir/test/Dialect/Standard/canonicalize.mlir

406–410

You are right!

[mehdi_amini]

arc diff --verbatim updates the summary on Phab FYI

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

def IndexCastOp : ArithmeticCastOp<"index_cast">, Arguments<(ins AnyType:$in)> {
   let summary = "cast between index and integer types";
   let description = [{
     Casts between integer scalars and 'index' scalars. Index is an integer of
     platform-specific bit width. If casting to a wider integer, the value is
     sign-extended. If casting to a narrower integer, the value is truncated.
   }];
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MulFOp
 //===----------------------------------------------------------------------===//
 
 def MulFOp : FloatBinaryOp<"mulf"> {
   let summary = "floating point multiplication operation";
let builders = [
     }]>];
 
   let parser = [{
     return impl::parseCastOp(parser, result);
   }];
   let printer = [{
     return printStandardCastOp(this->getOperation(), p);
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // SIToFPOp
 //===----------------------------------------------------------------------===//
 
 def SIToFPOp : ArithmeticCastOp<"sitofp">, Arguments<(ins AnyType:$in)> {
   let summary = "cast from integer type to floating-point";

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

OpFoldResult IndexCastOp::fold(ArrayRef<Attribute> cstOperands) {
   // A little hack because we go through int.  Otherwise, the size
   // of the constant might need to change.
   if (auto value = cstOperands[0].dyn_cast_or_null<IntegerAttr>())
     return IntegerAttr::get(getType(), value.getInt());
 
   return {};
 }
 
+namespace {
+///  index_cast(sign_extend x) => index_cast(x)
+struct IndexCastOfSExt : public OpRewritePattern<IndexCastOp> {
+  using OpRewritePattern<IndexCastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(IndexCastOp op,
+                                PatternRewriter &rewriter) const override {
+
+    if (auto extop = op.getOperand().getDefiningOp<SignExtendIOp>()) {
+      op.setOperand(extop.getOperand());
+      return success();
+    }
+    return failure();
+  }
+};
+
+} // namespace
+
+void IndexCastOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                              MLIRContext *context) {
+  results.insert<IndexCastOfSExt>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // MulFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult MulFOp::fold(ArrayRef<Attribute> operands) {
   return constFoldBinaryOp<FloatAttr>(
       operands, [](APFloat a, APFloat b) { return a * b; });
 }
static LogicalResult verify(SignExtendIOp op) {
   if (srcType.cast<IntegerType>().getWidth() >=
       dstType.cast<IntegerType>().getWidth())
     return op.emitError("result type ")
            << dstType << " must be wider than operand type " << srcType;
 
   return success();
 }
 
+OpFoldResult SignExtendIOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.size() == 1 && "unary operation takes one operand");
+
+  if (!operands[0])
+    return {};
+
+  if (auto lhs = operands[0].dyn_cast<IntegerAttr>()) {
+    return IntegerAttr::get(
+        getType(), lhs.getValue().sext(getType().getIntOrFloatBitWidth()));
+  }
+
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // SignedDivIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult SignedDivIOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 2 && "binary operation takes two operands");
 
   // Don't fold if it would overflow or if it requires a division by zero.

 
 OpFoldResult TruncateIOp::fold(ArrayRef<Attribute> operands) {
   // trunci(zexti(a)) -> a
   // trunci(sexti(a)) -> a
   if (matchPattern(getOperand(), m_Op<ZeroExtendIOp>()) ||
       matchPattern(getOperand(), m_Op<SignExtendIOp>()))
     return getOperand().getDefiningOp()->getOperand(0);
 
-  return nullptr;
+  assert(operands.size() == 1 && "unary operation takes one operand");
+
+  if (!operands[0])
+    return {};
+
+  if (auto lhs = operands[0].dyn_cast<IntegerAttr>()) {
+
+    return IntegerAttr::get(
+        getType(), lhs.getValue().trunc(getType().getIntOrFloatBitWidth()));
+  }
+
+  return {};
 }
 
 //===----------------------------------------------------------------------===//
 // UnsignedDivIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult UnsignedDivIOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 2 && "binary operation takes two operands");

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

// CHECK-NEXT: = llvm.fptrunc {{.*}} : vector<2xf64> to vector<2xf16>
   %1 = fptrunc %arg1: vector<2xf64> to vector<2xf16>
 // CHECK-NEXT: = llvm.fptrunc {{.*}} : vector<2xf64> to vector<2xf32>
   %2 = fptrunc %arg1: vector<2xf64> to vector<2xf32>
   return
 }
 
 // Check sign and zero extension and truncation of integers.
 // CHECK-LABEL: @integer_extension_and_truncation
-func @integer_extension_and_truncation() {
-// CHECK-NEXT:  %0 = llvm.mlir.constant(-3 : i3) : i3
-  %0 = constant 5 : i3
-// CHECK-NEXT: = llvm.sext %0 : i3 to i6
-  %1 = sexti %0 : i3 to i6
-// CHECK-NEXT: = llvm.zext %0 : i3 to i6
-  %2 = zexti %0 : i3 to i6
-// CHECK-NEXT: = llvm.trunc %0 : i3 to i2
-   %3 = trunci %0 : i3 to i2
+func @integer_extension_and_truncation(%arg0 : i3) {
+// CHECK-NEXT: = llvm.sext %arg0 : i3 to i6
+  %0 = sexti %arg0 : i3 to i6
+// CHECK-NEXT: = llvm.zext %arg0 : i3 to i6
+  %1 = zexti %arg0 : i3 to i6
+// CHECK-NEXT: = llvm.trunc %arg0 : i3 to i2
+   %2 = trunci %arg0 : i3 to i2
   return
 }
 
 // CHECK-LABEL: @dfs_block_order
 func @dfs_block_order(%arg0: i32) -> (i32) {
 // CHECK-NEXT:  %[[CST:.*]] = llvm.mlir.constant(42 : i32) : i32
   %0 = constant 42 : i32
 // CHECK-NEXT:  llvm.br ^bb2

mlir/test/Dialect/Standard/canonicalize.mlir

 
 // CHECK-LABEL: @select_cmp_ne_select
 //       CHECK:   return %arg0
 func @select_cmp_ne_select(%arg0: i64, %arg1: i64) -> i64 {
   %0 = cmpi ne, %arg0, %arg1 : i64
   %1 = select %0, %arg0, %arg1 : i64
   return %1 : i64
 }
+
+// -----
+
+// CHECK-LABEL: @indexCastOfSignExtend
+//       CHECK:   %[[res:.+]] = index_cast %arg0 : i8 to index
+//       CHECK:   return %[[res]]
+func @indexCastOfSignExtend(%arg0: i8) -> index {
+  %ext = sexti %arg0 : i8 to i16
+  %idx = index_cast %ext : i16 to index
+  return %idx : index
+}
+
+// CHECK-LABEL: @signExtendConstant
+//       CHECK:   %[[cres:.+]] = constant -2 : i16
+//       CHECK:   return %[[cres]]
+func @signExtendConstant() -> i16 {
+  %c-2 = constant -2 : i8
+  %ext = sexti %c-2 : i8 to i16
+  return %ext : i16
+}
+
+// CHECK-LABEL: @truncConstant
+//       CHECK:   %[[cres:.+]] = constant -2 : i16
+//       CHECK:   return %[[cres]]
+func @truncConstant(%arg0: i8) -> i16 {
+  %c-2 = constant -2 : i32
+  %tr = trunci %c-2 : i32 to i16
+  return %tr : i16
+}