[Mlir-commits] [mlir] [mlir][math] Fix APInt assertion in IPowI folding for i1 types (PR #177874)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Miloš Poletanović (milos1397)

<details>
<summary>Changes</summary>

Fixes #<!-- -->177831 
I have added a guard clause at the beginning of the fold method lambda to skip folding if the bitwidth of the operands is less than 2. This prevents the folder from attempting to create an invalid `oneValue` for i1 types.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Math/IR/MathOps.cpp (+2) 
- (added) mlir/test/Transforms/remove-dead-values-one-bit-width.mlir (+13) 


``````````diff
diff --git a/mlir/lib/Dialect/Math/IR/MathOps.cpp b/mlir/lib/Dialect/Math/IR/MathOps.cpp
index bbeef0f6ee9e5..2e99160a22d4c 100644
--- a/mlir/lib/Dialect/Math/IR/MathOps.cpp
+++ b/mlir/lib/Dialect/Math/IR/MathOps.cpp
@@ -369,6 +369,8 @@ OpFoldResult math::IPowIOp::fold(FoldAdaptor adaptor) {
       adaptor.getOperands(),
       [](const APInt &base, const APInt &power) -> std::optional<APInt> {
         unsigned width = base.getBitWidth();
+        if (width < 2)
+          return {};
         auto zeroValue = APInt::getZero(width);
         APInt oneValue{width, 1ULL, /*isSigned=*/true};
         APInt minusOneValue{width, -1ULL, /*isSigned=*/true};
diff --git a/mlir/test/Transforms/remove-dead-values-one-bit-width.mlir b/mlir/test/Transforms/remove-dead-values-one-bit-width.mlir
new file mode 100644
index 0000000000000..d9140e94f0bdd
--- /dev/null
+++ b/mlir/test/Transforms/remove-dead-values-one-bit-width.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt --remove-dead-values %s -o - | FileCheck %s
+// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
+
+// CHECK-LABEL:   func.func @main() {
+// CHECK:           return
+// CHECK:         }
+func.func @main() {
+  %1 = vector.constant_mask [0, 0] : vector<2x4xi1>
+  %2 = arith.constant dense<false> : vector<2x4xi1>
+  %5 = math.ipowi %1, %2 : vector<2x4xi1>
+  return
+}
+

``````````

</details>


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