[Mlir-commits] [mlir] [mlir][arith] Fix issue where setOperand is called inside the fold fu… (PR #204506)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: donald chen (cxy-1993)

<details>
<summary>Changes</summary>

…nction

---
Full diff: https://github.com/llvm/llvm-project/pull/204506.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/Arith/IR/ArithOps.cpp (+62-22) 


``````````diff
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index dc3887c3e0b0e..92f4e519f51a4 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -1735,26 +1735,12 @@ OpFoldResult arith::TruncIOp::fold(FoldAdaptor adaptor) {
     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 (llvm::cast<IntegerType>(srcType).getWidth() >
-        llvm::cast<IntegerType>(dstType).getWidth()) {
-      setOperand(src);
-      return getResult();
-    }
-
     // trunci(zexti(a)) -> a
     // trunci(sexti(a)) -> a
     if (srcType == dstType)
       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());
   unsigned bitWidth = llvm::cast<IntegerType>(resType).getWidth();
   return constFoldCastOp<IntegerAttr, IntegerAttr>(
@@ -1764,6 +1750,38 @@ OpFoldResult arith::TruncIOp::fold(FoldAdaptor adaptor) {
       });
 }
 
+struct TruncIOfExtI : public OpRewritePattern<arith::TruncIOp> {
+  using OpRewritePattern<arith::TruncIOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(arith::TruncIOp op,
+                                PatternRewriter &rewriter) const override {
+    if (matchPattern(op.getOperand(), m_Op<arith::ExtUIOp>()) ||
+        matchPattern(op.getOperand(), m_Op<arith::ExtSIOp>())) {
+      Value src = op.getOperand().getDefiningOp()->getOperand(0);
+      Type srcType = getElementTypeOrSelf(src.getType());
+      Type dstType = getElementTypeOrSelf(op.getType());
+      // trunci(zexti(a)) -> trunci(a)
+      // trunci(sexti(a)) -> trunci(a)
+      if (llvm::cast<IntegerType>(srcType).getWidth() >
+          llvm::cast<IntegerType>(dstType).getWidth()) {
+        rewriter.modifyOpInPlace(op, [&](){
+          op.setOperand(src);
+        });
+        return success();
+      }
+    }
+
+    // trunci(trunci(a)) -> trunci(a))
+    if (matchPattern(op.getOperand(), m_Op<arith::TruncIOp>())) {
+      rewriter.modifyOpInPlace(op, [&](){
+        op.setOperand(op.getOperand().getDefiningOp()->getOperand(0));
+      });
+      return success();
+    }
+    return failure();
+  }
+};
+
 bool arith::TruncIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkWidthChangeCast<std::less, IntegerType>(inputs, outputs);
 }
@@ -1771,7 +1789,7 @@ bool arith::TruncIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
 void arith::TruncIOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                   MLIRContext *context) {
   patterns
-      .add<TruncIExtSIToExtSI, TruncIExtUIToExtUI, TruncIShrSIToTrunciShrUI>(
+      .add<TruncIExtSIToExtSI, TruncIExtUIToExtUI, TruncIShrSIToTrunciShrUI, TruncIOfExtI>(
           context);
 }
 
@@ -1796,12 +1814,6 @@ OpFoldResult arith::TruncFOp::fold(FoldAdaptor adaptor) {
     if (llvm::APFloatBase::isRepresentableBy(
             srcType.getFloatSemantics(),
             intermediateType.getFloatSemantics())) {
-      // truncf(extf(a)) -> truncf(a)
-      if (srcType.getWidth() > resElemType.getWidth()) {
-        setOperand(src);
-        return getResult();
-      }
-
       // truncf(extf(a)) -> a
       if (srcType == resElemType)
         return src;
@@ -1824,9 +1836,37 @@ OpFoldResult arith::TruncFOp::fold(FoldAdaptor adaptor) {
       });
 }
 
+struct TruncFOfExtF : public OpRewritePattern<arith::TruncFOp> {
+  using OpRewritePattern<arith::TruncFOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(arith::TruncFOp op,
+                                PatternRewriter &rewriter) const override {
+    auto resElemType = cast<FloatType>(getElementTypeOrSelf(op.getType()));
+    if (auto extOp = op.getOperand().getDefiningOp<arith::ExtFOp>()) {
+      Value src = extOp.getIn();
+      auto srcType = cast<FloatType>(getElementTypeOrSelf(src.getType()));
+      auto intermediateType =
+          cast<FloatType>(getElementTypeOrSelf(extOp.getType()));
+      // Check if the srcType is representable in the intermediateType.
+      if (llvm::APFloatBase::isRepresentableBy(
+              srcType.getFloatSemantics(),
+              intermediateType.getFloatSemantics())) {
+        // truncf(extf(a)) -> truncf(a)
+        if (srcType.getWidth() > resElemType.getWidth()) {
+          rewriter.modifyOpInPlace(op, [&](){
+            op.setOperand(src);
+          });
+          return success();
+        }
+      }
+    }
+    return failure();
+  }
+};
+
 void arith::TruncFOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                   MLIRContext *context) {
-  patterns.add<TruncFSIToFPToSIToFP, TruncFUIToFPToUIToFP>(context);
+  patterns.add<TruncFSIToFPToSIToFP, TruncFUIToFPToUIToFP, TruncFOfExtF>(context);
 }
 
 bool arith::TruncFOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {

``````````

</details>


https://github.com/llvm/llvm-project/pull/204506