[Mlir-commits] [mlir] [mlir][SPIR-V] Add folder for spirv.SMulExtended(x, 1) (PR #214137)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-spirv

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>



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


2 Files Affected:

- (modified) mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp (+13-10) 
- (modified) mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir (+12) 


``````````diff
diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
index 2d5c4d7d3fd0e..ac6704cc51959 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
@@ -246,23 +246,18 @@ struct MulExtendedFold final : OpRewritePattern<MulOp> {
   }
 };
 
-using SMulExtendedOpFold = MulExtendedFold<spirv::SMulExtendedOp, true>;
-void spirv::SMulExtendedOp::getCanonicalizationPatterns(
-    RewritePatternSet &patterns, MLIRContext *context) {
-  patterns.add<SMulExtendedOpFold>(context);
-}
-
-struct UMulExtendedOpXOne final : OpRewritePattern<spirv::UMulExtendedOp> {
-  using Base::Base;
+template <typename MulOp>
+struct MulExtendedOpXOne final : OpRewritePattern<MulOp> {
+  using OpRewritePattern<MulOp>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(spirv::UMulExtendedOp op,
+  LogicalResult matchAndRewrite(MulOp op,
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     Value lhs = op.getOperand1();
     Value rhs = op.getOperand2();
     Type constituentType = lhs.getType();
 
-    // umulextended (x, 1) = <x, 0>
+    // [su]mulextended (x, 1) = <x, 0>
     if (matchPattern(rhs, m_One())) {
       Value zero = spirv::ConstantOp::getZero(constituentType, loc, rewriter);
       Value constituents[2] = {lhs, zero};
@@ -275,7 +270,15 @@ struct UMulExtendedOpXOne final : OpRewritePattern<spirv::UMulExtendedOp> {
   }
 };
 
+using SMulExtendedOpFold = MulExtendedFold<spirv::SMulExtendedOp, true>;
+using SMulExtendedOpXOne = MulExtendedOpXOne<spirv::SMulExtendedOp>;
+void spirv::SMulExtendedOp::getCanonicalizationPatterns(
+    RewritePatternSet &patterns, MLIRContext *context) {
+  patterns.add<SMulExtendedOpFold, SMulExtendedOpXOne>(context);
+}
+
 using UMulExtendedOpFold = MulExtendedFold<spirv::UMulExtendedOp, false>;
+using UMulExtendedOpXOne = MulExtendedOpXOne<spirv::UMulExtendedOp>;
 void spirv::UMulExtendedOp::getCanonicalizationPatterns(
     RewritePatternSet &patterns, MLIRContext *context) {
   patterns.add<UMulExtendedOpFold, UMulExtendedOpXOne>(context);
diff --git a/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir b/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
index 94d9c53db0bbc..713471821b4bc 100644
--- a/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
+++ b/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
@@ -502,6 +502,18 @@ func.func @smulextended_x_0(%arg0 : i32) -> !spirv.struct<(i32, i32)> {
   return %0 : !spirv.struct<(i32, i32)>
 }
 
+// CHECK-LABEL: @smulextended_x_1
+// CHECK-SAME: (%[[ARG:.*]]: i32)
+func.func @smulextended_x_1(%arg0 : i32) -> !spirv.struct<(i32, i32)> {
+  // CHECK: %[[C0:.*]] = spirv.Constant 0
+  // CHECK: %[[RET:.*]] = spirv.CompositeConstruct %[[ARG]], %[[C0]]
+  %c0 = spirv.Constant 1 : i32
+  %0 = spirv.SMulExtended %arg0, %c0 : !spirv.struct<(i32, i32)>
+
+  // CHECK: return %[[RET]]
+  return %0 : !spirv.struct<(i32, i32)>
+}
+
 // CHECK-LABEL: @const_fold_scalar_smulextended
 func.func @const_fold_scalar_smulextended() -> (!spirv.struct<(i32, i32)>, !spirv.struct<(i32, i32)>) {
   %c5 = spirv.Constant 5 : i32

``````````

</details>


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