[MLIR][SPIRV] Add rewrite pattern to convert select+cmp into GLSL clamp.

[ergawy]

Summary

Adds rewrite patterns to convert select+cmp instructions into clamp
instructions whenever possible. Support is added to convert:

• FOrdLessThan, FOrdLessThanEqual to GLSLFClampOp.
• SLessThan, SLessThanEqual to GLSLSClampOp.
• ULessThan, ULessThanEqual to GLSLUClampOp.

[ergawy]

2 things I still need to figure out/need to make sure I got right:

• How to enable these patterns only when the GLSL.std.450 extension is enabled?
• Whether you think it's suitable to consider these rewrites as canonicalization patterns or we should move them somewhere else?

[mravishankar]

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

239

Nice! but I am not sure we want to add it for general canonicalizers. GLSL is an extension and you dont necessarily need to use it. Maybe put these patterns into another method

populateSPIRVGLSLCanonicalizationPatterns(...)

Then its upto client to use that. For testing this just create a simple pass that just runs those canonicalization patterns

[ergawy]

• Split GLSL-specific patterns in their own file.
• Add a test pass to test them.

[ergawy]

Fix file organization issue.

[ergawy]

Remove empty line.

[ergawy]

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

239

Thanks @mravishankar for the review. I re-organized the code according to your suggestion. Let me know if further changes need to be done here.

[ergawy]

Add missing header guard.

[ergawy]

Handle some minor clang-tidy comments.

[mravishankar]

Thanks! LGTM

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

def SPV_SelectOp : SPV_Op<"Select",
   );
 
   let builders = [
     OpBuilderDAG<(ins "Value":$cond, "Value":$trueValue, "Value":$falseValue)>];
 
   let assemblyFormat = [{
     operands attr-dict `:` type($condition) `,` type($result)
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 // -----
 
 def SPV_UGreaterThanOp : SPV_LogicalBinaryOp<"UGreaterThan", SPV_Integer, []> {
   let summary = [{
     Unsigned-integer comparison if Operand 1 is greater than  Operand 2.
   }];

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

 
 void spirv::LogicalNotOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<ConvertLogicalNotOfIEqual, ConvertLogicalNotOfINotEqual,
                  ConvertLogicalNotOfLogicalEqual,
                  ConvertLogicalNotOfLogicalNotEqual>(context);
 }
 
+void spirv::SelectOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<ConvertComparisonIntoClampSPV_FOrdLessThanOp,
+                 ConvertComparisonIntoClampSPV_FOrdLessThanEqualOp,
+                 ConvertComparisonIntoClampSPV_SLessThanOp,
+                 ConvertComparisonIntoClampSPV_SLessThanEqualOp,
+                 ConvertComparisonIntoClampSPV_ULessThanOp,
+                 ConvertComparisonIntoClampSPV_ULessThanEqualOp>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // spv.LogicalOr
 //===----------------------------------------------------------------------===//
 
 OpFoldResult spirv::LogicalOrOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 2 && "spv.LogicalOr should take two operands");
 
   if (auto rhs = getScalarOrSplatBoolAttr(operands.back())) {

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

 
 def ConvertLogicalNotOfLogicalEqual : Pat<
     (SPV_LogicalNotOp (SPV_LogicalEqualOp $lhs, $rhs)),
     (SPV_LogicalNotEqualOp $lhs, $rhs)>;
 
 def ConvertLogicalNotOfLogicalNotEqual : Pat<
     (SPV_LogicalNotOp (SPV_LogicalNotEqualOp $lhs, $rhs)),
     (SPV_LogicalEqualOp $lhs, $rhs)>;
+
+//===----------------------------------------------------------------------===//
+// Re-write spv.Select + spv.<less_than_op> to a suitable variant of
+// spv.<glsl_clamp_op>
+//===----------------------------------------------------------------------===//
+
+def ValuesAreEqual : Constraint<CPred<"$0 == $1">>;
+
+foreach CmpClampPair = [
+    [SPV_FOrdLessThanOp, SPV_GLSLFClampOp],
+    [SPV_FOrdLessThanEqualOp, SPV_GLSLFClampOp],
+    [SPV_SLessThanOp, SPV_GLSLSClampOp],
+    [SPV_SLessThanEqualOp, SPV_GLSLSClampOp],
+    [SPV_ULessThanOp, SPV_GLSLUClampOp],
+    [SPV_ULessThanEqualOp, SPV_GLSLUClampOp]] in {
+def ConvertComparisonIntoClamp#CmpClampPair[0] : Pat<
+    (SPV_SelectOp
+        (CmpClampPair[0]
+            (SPV_SelectOp:$middle0
+                (CmpClampPair[0] $min, $input),
+                $input,
+                $min
+            ),
+            $max
+        ),
+        $middle1,
+        $max),
+    (CmpClampPair[1] $input, $min, $max),
+    [(ValuesAreEqual $middle0, $middle1)]>;
+}

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

^else:
     spv.Store "Function" %3, %2 ["Aligned", 8] : vector<3xi32>
     spv.Branch ^merge
 
   ^merge:
     spv.mlir.merge
   }
   spv.Return
 }
+
+// -----
+
+// CHECK: func @clamp_fordlessthan(%[[INPUT:.*]]: f32)
+func @clamp_fordlessthan(%input: f32) -> f32{
+  // CHECK: %[[MIN:.*]] = spv.constant
+  %min = spv.constant 0.5 : f32
+  // CHECK: %[[MAX:.*]] = spv.constant
+  %max = spv.constant 1.0 : f32
+
+  // CHECK: [[RES:%.*]] = spv.GLSL.FClamp %[[INPUT]], %[[MIN]], %[[MAX]]
+  %0 = spv.FOrdLessThan %min, %input : f32
+  %mid = spv.Select %0, %input, %min : i1, f32
+  %1 = spv.FOrdLessThan %mid, %max : f32
+  %2 = spv.Select %1, %mid, %max : i1, f32
+
+  // CHECK-NEXT: spv.ReturnValue [[RES]]
+  spv.ReturnValue %2 : f32
+}
+
+// -----
+
+// CHECK: func @clamp_fordlessthanequal(%[[INPUT:.*]]: f32)
+func @clamp_fordlessthanequal(%input: f32) -> f32 {
+  // CHECK: %[[MIN:.*]] = spv.constant
+  %min = spv.constant 0.5 : f32
+  // CHECK: %[[MAX:.*]] = spv.constant
+  %max = spv.constant 1.0 : f32
+
+  // CHECK: [[RES:%.*]] = spv.GLSL.FClamp %[[INPUT]], %[[MIN]], %[[MAX]]
+  %0 = spv.FOrdLessThanEqual %min, %input : f32
+  %mid = spv.Select %0, %input, %min : i1, f32
+  %1 = spv.FOrdLessThanEqual %mid, %max : f32
+  %2 = spv.Select %1, %mid, %max : i1, f32
+
+  // CHECK-NEXT: spv.ReturnValue [[RES]]
+  spv.ReturnValue %2 : f32
+}
+
+// -----
+
+// CHECK: func @clamp_slessthan(%[[INPUT:.*]]: si32)
+func @clamp_slessthan(%input: si32) -> si32 {
+  // CHECK: %[[MIN:.*]] = spv.constant
+  %min = spv.constant 0 : si32
+  // CHECK: %[[MAX:.*]] = spv.constant
+  %max = spv.constant 10 : si32
+
+  // CHECK: [[RES:%.*]] = spv.GLSL.SClamp %[[INPUT]], %[[MIN]], %[[MAX]]
+  %0 = spv.SLessThan %min, %input : si32
+  %mid = spv.Select %0, %input, %min : i1, si32
+  %1 = spv.SLessThan %mid, %max : si32
+  %2 = spv.Select %1, %mid, %max : i1, si32
+
+  // CHECK-NEXT: spv.ReturnValue [[RES]]
+  spv.ReturnValue %2 : si32
+}
+
+// -----
+
+// CHECK: func @clamp_slessthanequal(%[[INPUT:.*]]: si32)
+func @clamp_slessthanequal(%input: si32) -> si32 {
+  // CHECK: %[[MIN:.*]] = spv.constant
+  %min = spv.constant 0 : si32
+  // CHECK: %[[MAX:.*]] = spv.constant
+  %max = spv.constant 10 : si32
+
+  // CHECK: [[RES:%.*]] = spv.GLSL.SClamp %[[INPUT]], %[[MIN]], %[[MAX]]
+  %0 = spv.SLessThanEqual %min, %input : si32
+  %mid = spv.Select %0, %input, %min : i1, si32
+  %1 = spv.SLessThanEqual %mid, %max : si32
+  %2 = spv.Select %1, %mid, %max : i1, si32
+
+  // CHECK-NEXT: spv.ReturnValue [[RES]]
+  spv.ReturnValue %2 : si32
+}
+
+// -----
+
+// CHECK: func @clamp_ulessthan(%[[INPUT:.*]]: i32)
+func @clamp_ulessthan(%input: i32) -> i32 {
+  // CHECK: %[[MIN:.*]] = spv.constant
+  %min = spv.constant 0 : i32
+  // CHECK: %[[MAX:.*]] = spv.constant
+  %max = spv.constant 10 : i32
+
+  // CHECK: [[RES:%.*]] = spv.GLSL.UClamp %[[INPUT]], %[[MIN]], %[[MAX]]
+  %0 = spv.ULessThan %min, %input : i32
+  %mid = spv.Select %0, %input, %min : i1, i32
+  %1 = spv.ULessThan %mid, %max : i32
+  %2 = spv.Select %1, %mid, %max : i1, i32
+
+  // CHECK-NEXT: spv.ReturnValue [[RES]]
+  spv.ReturnValue %2 : i32
+}
+
+// -----
+
+// CHECK: func @clamp_ulessthanequal(%[[INPUT:.*]]: i32)
+func @clamp_ulessthanequal(%input: i32) -> i32 {
+  // CHECK: %[[MIN:.*]] = spv.constant
+  %min = spv.constant 0 : i32
+  // CHECK: %[[MAX:.*]] = spv.constant
+  %max = spv.constant 10 : i32
+
+  // CHECK: [[RES:%.*]] = spv.GLSL.UClamp %[[INPUT]], %[[MIN]], %[[MAX]]
+  %0 = spv.ULessThanEqual %min, %input : i32
+  %mid = spv.Select %0, %input, %min : i1, i32
+  %1 = spv.ULessThanEqual %mid, %max : i32
+  %2 = spv.Select %1, %mid, %max : i1, i32
+
+  // CHECK-NEXT: spv.ReturnValue [[RES]]
+  spv.ReturnValue %2 : i32
+}