[Mlir-commits] [mlir] [mlir][Utils] Add overflow<nsw> flags to dimension calculations in ReshapeOpsUtils (PR #178766)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Zhewen Yu (Yu-Zhewen)

<details>
<summary>Changes</summary>

This follows the convention that `tensor/memref` indexing calculations generally should never signed overflow (https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Arith/Utils/Utils.h#L123-L125), and enables CSE to properly merge duplicate dimension computations.

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


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h (+4-2) 
- (modified) mlir/test/Dialect/Tensor/canonicalize.mlir (+1-1) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
index 64c125024d906..9345b0b167d27 100644
--- a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
@@ -332,11 +332,13 @@ struct ComposeCollapseOfExpandOp : public OpRewritePattern<CollapseOpTy> {
         // the first dynamic size.
         Value result = dynamicSizes[0];
         for (Value v : llvm::drop_begin(dynamicSizes))
-          result = arith::MulIOp::create(rewriter, loc, result, v);
+          result = arith::MulIOp::create(rewriter, loc, result, v,
+                                         arith::IntegerOverflowFlags::nsw);
         if (numStaticElems != 1) {
           result = arith::MulIOp::create(
               rewriter, loc, result,
-              arith::ConstantIndexOp::create(rewriter, loc, numStaticElems));
+              arith::ConstantIndexOp::create(rewriter, loc, numStaticElems),
+              arith::IntegerOverflowFlags::nsw);
         }
         newOutputShape.push_back(result);
       }
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 95c5b8c91edf5..03b81adae5d71 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -1264,7 +1264,7 @@ func.func @compose_collapse_of_expand_partially_dynamic(%arg0: tensor<?xf16>, %a
 //  CHECK-SAME:   %[[ORIG_D2:.[a-zA-Z0-9]+]]
 //  CHECK-SAME:   %[[ORIG_D3:.[a-zA-Z0-9]+]]
 //   CHECK-DAG:   %[[C32:.+]] = arith.constant 32
-//       CHECK:   %[[COLLAPSED_D2:.+]] = arith.muli %[[ORIG_D3]], %[[C32]]
+//       CHECK:   %[[COLLAPSED_D2:.+]] = arith.muli %[[ORIG_D3]], %[[C32]] overflow<nsw>
 //       CHECK:   %[[RESULT:.+]] = tensor.expand_shape %[[SRC]]
 //  CHECK-SAME:     [0, 1, 2]
 //  CHECK-SAME:     output_shape [8, %[[ORIG_D2]], %[[COLLAPSED_D2]]]

``````````

</details>


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