[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 ↗

(On Diff #313032)

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 ↗

(On Diff #313032)

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/SPIRVGLSLCanonicalization.h

+//===- SPIRVGLSLCanonicalization.h - GLSL-specific patterns -----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares a function to register SPIR-V GLSL-specific
+// canonicalization patterns.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPIRV_IR_GLSLCANONICALIZATION_H_
+#define MLIR_DIALECT_SPIRV_IR_GLSLCANONICALIZATION_H_
+
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/PatternMatch.h"
+
+//===----------------------------------------------------------------------===//
+// GLSL canonicalization patterns
+//===----------------------------------------------------------------------===//
+
+namespace mlir {
+namespace spirv {
+void populateSPIRVGLSLCanonicalizationPatterns(
+    mlir::OwningRewritePatternList &results, mlir::MLIRContext *context);
+} // namespace spirv
+} // namespace mlir
+
+#endif // MLIR_DIALECT_SPIRV_IR_GLSLCANONICALIZATION_H_

mlir/lib/Dialect/SPIRV/IR/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS SPIRVCanonicalization.td)
 mlir_tablegen(SPIRVCanonicalization.inc -gen-rewriters)
 add_public_tablegen_target(MLIRSPIRVCanonicalizationIncGen)
 
 add_mlir_dialect_library(MLIRSPIRV
   SPIRVAttributes.cpp
   SPIRVCanonicalization.cpp
+  SPIRVGLSLCanonicalization.cpp
   SPIRVDialect.cpp
   SPIRVEnums.cpp
   SPIRVOps.cpp
   SPIRVTypes.cpp
   TargetAndABI.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SPIRV

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/lib/Dialect/SPIRV/IR/SPIRVGLSLCanonicalization.cpp

+//===- SPIRVGLSLCanonicalization.cpp - SPIR-V GLSL canonicalization patterns =//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the canonicalization patterns for SPIR-V GLSL-specific ops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SPIRV/IR/SPIRVGLSLCanonicalization.h"
+
+#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
+
+using namespace mlir;
+
+namespace {
+#include "SPIRVCanonicalization.inc"
+}
+
+namespace mlir {
+namespace spirv {
+void populateSPIRVGLSLCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<ConvertComparisonIntoClampSPV_FOrdLessThanOp,
+                 ConvertComparisonIntoClampSPV_FOrdLessThanEqualOp,
+                 ConvertComparisonIntoClampSPV_SLessThanOp,
+                 ConvertComparisonIntoClampSPV_SLessThanEqualOp,
+                 ConvertComparisonIntoClampSPV_ULessThanOp,
+                 ConvertComparisonIntoClampSPV_ULessThanEqualOp>(context);
+}
+} // namespace spirv
+} // namespace mlir

mlir/test/Dialect/SPIRV/Transforms/glsl_canonicalize.mlir

+// RUN: mlir-opt -test-spirv-glsl-canonicalization -split-input-file -verify-diagnostics %s | FileCheck %s
+
+// 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
+}

mlir/test/lib/Dialect/SPIRV/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRSPIRVTestPasses
   TestAvailability.cpp
   TestEntryPointAbi.cpp
+  TestGLSLCanonicalization.cpp
   TestModuleCombiner.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SPIRV
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
 

mlir/test/lib/Dialect/SPIRV/TestGLSLCanonicalization.cpp

+//===- TestGLSLCanonicalization.cpp - Pass to test GLSL-specific pattterns ===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SPIRV/IR/SPIRVGLSLCanonicalization.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVModule.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+class TestGLSLCanonicalizationPass
+    : public PassWrapper<TestGLSLCanonicalizationPass,
+                         OperationPass<mlir::ModuleOp>> {
+public:
+  TestGLSLCanonicalizationPass() = default;
+  TestGLSLCanonicalizationPass(const TestGLSLCanonicalizationPass &) {}
+  void runOnOperation() override;
+};
+} // namespace
+
+void TestGLSLCanonicalizationPass::runOnOperation() {
+  OwningRewritePatternList patterns;
+  spirv::populateSPIRVGLSLCanonicalizationPatterns(patterns, &getContext());
+  applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+}
+
+namespace mlir {
+void registerTestSpirvGLSLCanonicalizationPass() {
+  PassRegistration<TestGLSLCanonicalizationPass> registration(
+      "test-spirv-glsl-canonicalization",
+      "Tests SPIR-V canonicalization patterns for GLSL extension.");
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestLoopPermutationPass();
 void registerTestMatchers();
 void registerTestPrintDefUsePass();
 void registerTestPrintNestingPass();
 void registerTestReducer();
 void registerTestSpirvEntryPointABIPass();
+void registerTestSpirvGLSLCanonicalizationPass();
 void registerTestSpirvModuleCombinerPass();
 void registerTestTraitsPass();
 void registerTosaTestQuantUtilAPIPass();
 void registerVectorizerTestPass();
 
 namespace test {
 void registerConvertCallOpPass();
 void registerInliner();
void registerTestPasses() {
   registerTestFunc();
   registerTestGpuMemoryPromotionPass();
   registerTestLoopPermutationPass();
   registerTestMatchers();
   registerTestPrintDefUsePass();
   registerTestPrintNestingPass();
   registerTestReducer();
   registerTestSpirvEntryPointABIPass();
+  registerTestSpirvGLSLCanonicalizationPass();
   registerTestSpirvModuleCombinerPass();
   registerTestTraitsPass();
   registerVectorizerTestPass();
   registerTosaTestQuantUtilAPIPass();
 
   test::registerConvertCallOpPass();
   test::registerInliner();
   test::registerMemRefBoundCheck();