[MLIR][SPIRV] Add a canonicalization pattern for UModOp

[nbpatel]

Summary

Add a transformation for a pattern like

%6 = "spirv.Constant"() <{value = 32 : i32}> : () -> i32
%7 = "spirv.UMod"(%5, %6) : (i32, i32) -> i32
%8 = "spirv.Constant"() <{value = 4 : i32}> : () -> i32
%9 = "spirv.UMod"(%7, %8) : (i32, i32) -> i32

to transform to

%6 = "spirv.Constant"() <{value = 32 : i32}> : () -> i32
%7 = "spirv.UMod"(%5, %6) : (i32, i32) -> i32
%8 = "spirv.Constant"() <{value = 4 : i32}> : () -> i32
%9= "spirv.UMod"(%5, %8) : (i32, i32) -> i32

[kuhar]

Thanks for coming up with this canon pattern, @nbpatel. I left some suggestions but overall this looks like a useful pattern to me. I have two high-level questions:

1. Have you tried defining it in DDR instead (SPIRVCanonicalization.td)? The implementation may be simpler and more concise there.

2. Do you see pattern show up a lot in your IR? What do you think about adding it on the level of the Arith dialect instead/in addition?

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

128

nit: transformation works only if --> the transformation is only applied when?
'works only if' could be interpreted as the transformation being broken or incomplete

130–132

nit: drop this repeated header

134–135

Could you update pattern to the same style as the other patterns in this file? EG:

struct CombineChainedAccessChain final
    : OpRewritePattern<spirv::AccessChainOp> {
  using OpRewritePattern::OpRewritePattern;

This is to maintain consistency.

139–140

We can combine these two:

auto prevUMod = umodOp.getOperand(0).getDefiningOp<spirv::UModOp>();

145–146

Match with m_Constant instead

147

same here

152–153

1. Member cast functions are deprecated, use free functions instead: https://mlir.llvm.org/deprecation/
2. Could the value be a vector type?

162–164

Use rewriter.replaceOpWithNewOp

mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir

483–484

Can we also check that the expected values were indeed returned? Also in the previous testcase

487–488

I think it would help to have 2 more tests:

1. For when the transformation does not apply because of the divisors not being multiples
2. Check that the pattern handles vectors of integers.

[nbpatel]

Address PR feedback.

Yes, we do see a lot of such code in our IR and we would want to clean it up a bit for easier readability and debug.

I haven't tried it with vector

[nbpatel]

@kuhar do we need to handle vectors in this PR? can it be an incremental PR?
I am not sure how complex/simple the change will be for handling vectors

[kuhar]

Thanks for all the fixes.

In D152341#4404039, @nbpatel wrote:

@kuhar do we need to handle vectors in this PR? can it be an incremental PR?
I am not sure how complex/simple the change will be for handling vectors

In the previous iteration the code would crash on vectors because of casts. Now I think it will bail out:

IntegerAttr prevValue, currValue;
if (!matchPattern(prevUMod.getOperand(1), m_Constant(&prevValue)) ||
    !matchPattern(umodOp.getOperand(1), m_Constant(&currValue)))
  return failure();

Supporting vectors should be easy, so I think it makes sense to handle it in this PR. You can find examples of similar canon patterns in ArithOps.cpp. If you decide not to do that, this deserves a testcase and a TODO at the very least.

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

139

nit: don't define multiple variables in the same statement

155–158

This should require only one new UMod op, no? In the example on top of the function only the second one is changed.

[nbpatel]

Made requested changes and added a TODO.

[kuhar]

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

130

Can you open a github issue, tag it with [mlir][spirv] and add the link inside TODO(link). This makes keeping track of unresolved issues much easier.

[kuhar]

Have you also seen something something like:
a = umod x, C1
b = umod a, C2
where C1 < C2?

I wonder if we should also support that and fold to b ==> a. This is separate but I'm curious if it shows up in your code.

For this PR, LGTM modulo the last comment. @antiagainst, could you also take a look?

[nbpatel]

Thanks @kuhar for reviewing and accepting the change. Can you also please help in merging it since I dont have permission to do so?

[nbpatel]

added the link for TODO

[kuhar]

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

130

[nbpatel]

No , haven't seen this pattern show up. But it would be a good add.

In D152341#4404308, @kuhar wrote:

Have you also seen something something like:
a = umod x, C1
b = umod a, C2
where C1 < C2?

I wonder if we should also support that and fold to b ==> a. This is separate but I'm curious if it shows up in your code.

For this PR, LGTM modulo the last comment. @antiagainst, could you also take a look?

[kuhar]

In D152341#4404313, @nbpatel wrote:

Thanks @kuhar for reviewing and accepting the change. Can you also please help in merging it since I dont have permission to do so?

Yeah I will do that once @antiagainst gives it a thumbs up.

[nbpatel]

Name: Nishant Patel
Email: [e-mail]

Information for merging.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td

let description = [{
     #### Example:
 
     ```mlir
     %4 = spirv.UMod %0, %1 : i32
     %5 = spirv.UMod %2, %3 : vector<4xi32>
 
     ```
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 #endif // MLIR_DIALECT_SPIRV_IR_ARITHMETIC_OPS

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

 } // namespace
 
 void spirv::AccessChainOp::getCanonicalizationPatterns(
     RewritePatternSet &results, MLIRContext *context) {
   results.add<CombineChainedAccessChain>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// spirv.UMod
+//===----------------------------------------------------------------------===//
+
+// Input:
+//    %0 = spirv.UMod %arg0, %const32 : i32
+//    %1 = spirv.UMod %0, %const4 : i32
+// Output:
+//    %0 = spirv.UMod %arg0, %const32 : i32
+//    %1 = spirv.UMod %arg0, %const4 : i32
+
+// The transformation is only applied if one divisor is a multiple of the other.
+
+// TODO(https://github.com/llvm/llvm-project/issues/63174): Add support for vector constants
+struct UModSimplification final : OpRewritePattern<spirv::UModOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spirv::UModOp umodOp,
+                                PatternRewriter &rewriter) const override {
+    auto prevUMod = umodOp.getOperand(0).getDefiningOp<spirv::UModOp>();
+    if (!prevUMod)
+      return failure();
+
+    IntegerAttr prevValue;
+    IntegerAttr currValue;
+    if (!matchPattern(prevUMod.getOperand(1), m_Constant(&prevValue)) ||
+        !matchPattern(umodOp.getOperand(1), m_Constant(&currValue)))
+      return failure();
+
+    APInt prevConstValue = prevValue.getValue();
+    APInt currConstValue = currValue.getValue();
+
+    // Ensure that one divisor is a multiple of the other. If not, fail the
+    // transformation.
+    if (prevConstValue.urem(currConstValue) != 0 &&
+        currConstValue.urem(prevConstValue) != 0)
+      return failure();
+
+    // The transformation is safe. Replace the existing UMod operation with a
+    // new UMod operation, using the original dividend and the current divisor.
+    rewriter.replaceOpWithNewOp<spirv::UModOp>(
+        umodOp, umodOp.getType(), prevUMod.getOperand(0), umodOp.getOperand(1));
+
+    return success();
+  }
+};
+
+void spirv::UModOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+                                                MLIRContext *context) {
+  patterns.insert<UModSimplification>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // spirv.BitcastOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult spirv::BitcastOp::fold(FoldAdaptor /*adaptor*/) {
   Value curInput = getOperand();
   if (getType() == curInput.getType())
     return curInput;
 

mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir

func.func @const_fold_vector_isub() -> vector<3xi32> {
   // CHECK: spirv.Constant dense<[45, -40, 99]>
   %0 = spirv.ISub %vc1, %vc2 : vector<3xi32>
   return %0: vector<3xi32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
+// spirv.UMod
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @umod_fold
+// CHECK-SAME: (%[[ARG:.*]]: i32)
+func.func @umod_fold(%arg0: i32) -> (i32, i32) {
+  // CHECK: %[[CONST4:.*]] = spirv.Constant 4
+  // CHECK: %[[CONST32:.*]] = spirv.Constant 32
+  %const1 = spirv.Constant 32 : i32
+  %0 = spirv.UMod %arg0, %const1 : i32
+  %const2 = spirv.Constant 4 : i32
+  %1 = spirv.UMod %0, %const2 : i32
+  // CHECK: %[[UMOD0:.*]] = spirv.UMod %[[ARG]], %[[CONST32]]
+  // CHECK: %[[UMOD1:.*]] = spirv.UMod %[[ARG]], %[[CONST4]]
+  // CHECK: return %[[UMOD0]], %[[UMOD1]]
+  return %0, %1: i32, i32
+}
+
+// CHECK-LABEL: @umod_fail_vector_fold
+// CHECK-SAME: (%[[ARG:.*]]: vector<4xi32>)
+func.func @umod_fail_vector_fold(%arg0: vector<4xi32>) -> (vector<4xi32>, vector<4xi32>) {
+  // CHECK: %[[CONST4:.*]] = spirv.Constant dense<4> : vector<4xi32>
+  // CHECK: %[[CONST32:.*]] = spirv.Constant dense<32> : vector<4xi32>
+  %const1 = spirv.Constant dense<32> : vector<4xi32>
+  %0 = spirv.UMod %arg0, %const1 : vector<4xi32>
+  // CHECK: %[[UMOD0:.*]] = spirv.UMod %[[ARG]], %[[CONST32]]
+  %const2 = spirv.Constant dense<4> : vector<4xi32>
+  %1 = spirv.UMod %0, %const2 : vector<4xi32>
+  // CHECK: %[[UMOD1:.*]] = spirv.UMod %[[UMOD0]], %[[CONST4]]
+  // CHECK: return %[[UMOD0]], %[[UMOD1]]
+  return %0, %1: vector<4xi32>, vector<4xi32>
+} 
+
+// CHECK-LABEL: @umod_fold_same_divisor
+// CHECK-SAME: (%[[ARG:.*]]: i32)
+func.func @umod_fold_same_divisor(%arg0: i32) -> (i32, i32) {
+  // CHECK: %[[CONST1:.*]] = spirv.Constant 32
+  %const1 = spirv.Constant 32 : i32
+  %0 = spirv.UMod %arg0, %const1 : i32
+  %const2 = spirv.Constant 32 : i32
+  %1 = spirv.UMod %0, %const2 : i32
+  // CHECK: %[[UMOD0:.*]] = spirv.UMod %[[ARG]], %[[CONST1]]
+  // CHECK: %[[UMOD1:.*]] = spirv.UMod %[[ARG]], %[[CONST1]]
+  // CHECK: return %[[UMOD0]], %[[UMOD1]]
+  return %0, %1: i32, i32
+}
+
+// CHECK-LABEL: @umod_fail_fold
+// CHECK-SAME: (%[[ARG:.*]]: i32)
+func.func @umod_fail_fold(%arg0: i32) -> (i32, i32) {
+  // CHECK: %[[CONST5:.*]] = spirv.Constant 5
+  // CHECK: %[[CONST32:.*]] = spirv.Constant 32
+  %const1 = spirv.Constant 32 : i32
+  %0 = spirv.UMod %arg0, %const1 : i32
+  // CHECK: %[[UMOD0:.*]] = spirv.UMod %[[ARG]], %[[CONST32]]
+  %const2 = spirv.Constant 5 : i32
+  %1 = spirv.UMod %0, %const2 : i32
+  // CHECK: %[[UMOD1:.*]] = spirv.UMod %[[UMOD0]], %[[CONST5]]
+  // CHECK: return %[[UMOD0]], %[[UMOD1]]
+  return %0, %1: i32, i32
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
 // spirv.LogicalAnd
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @convert_logical_and_true_false_scalar
 // CHECK-SAME: %[[ARG:.+]]: i1
 func.func @convert_logical_and_true_false_scalar(%arg: i1) -> (i1, i1) {
   %true = spirv.Constant true
   // CHECK: %[[FALSE:.+]] = spirv.Constant false