[mlir][spirv] Add some folders for spv.LogicalAnd/spv.LogicalOr

[antiagainst]

Summary

This commit handles folding spv.LogicalAnd/spv.LogicalOr when
one of the operands is constant true/false.

[rriddle]

mlir/lib/Dialect/SPIRV/SPIRVCanonicalization.cpp

41

Is the template here really necessary? Also, please use getSplatValue<bool>() instead.

[rriddle]

mlir/lib/Dialect/SPIRV/SPIRVCanonicalization.cpp

29

Would it be better to change this to something like:

static Optional<bool> getScalarOrSplatBoolAttr(Attribute boolAttr)

to avoid doing the same checks twice in each of the folders?

[rriddle]

mlir/lib/Dialect/SPIRV/SPIRVCanonicalization.cpp

27

The comment irVal seems to be out-of-date.

[antiagainst]

Oops. Pushed an older commit without squashing fixes...

[antiagainst]

Comments addressed as https://github.com/llvm/llvm-project/commit/63779fb462d828d16b87f427a6490dded842ca15

mlir/lib/Dialect/SPIRV/SPIRVCanonicalization.cpp

29

Good suggestion! Done.

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

let description = [{
 
     For example:
 
     ```
     %2 = spv.LogicalAnd %0, %1 : i1
     %2 = spv.LogicalAnd %0, %1 : vector<4xi1>
     ```
   }];
+
+  let hasFolder = 1;
 }
 
 // -----
 
 def SPV_LogicalEqualOp : SPV_LogicalBinaryOp<"LogicalEqual", SPV_Bool, [Commutative]> {
   let summary = [{
     Result is true if Operand 1 and Operand 2 have the same value. Result is
     false if Operand 1 and Operand 2 have different values.
let description = [{
 
     For example:
 
     ```
     %2 = spv.LogicalOr %0, %1 : i1
     %2 = spv.LogicalOr %0, %1 : vector<4xi1>
     ```
   }];
+
+  let hasFolder = 1;
 }
 
 // -----
 
 def SPV_SGreaterThanOp : SPV_LogicalBinaryOp<"SGreaterThan", SPV_Integer, []> {
   let summary = [{
     Signed-integer comparison if Operand 1 is greater than  Operand 2.
   }];

mlir/lib/Dialect/SPIRV/SPIRVCanonicalization.cpp

 #include "mlir/Support/Functional.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Common utility functions
 //===----------------------------------------------------------------------===//
 
+/// Returns true if the given `irVal` is a scalar or splat vector constant of
+/// the given `boolVal`.
+static bool isScalarOrSplatBoolAttr(Attribute boolAttr, bool boolVal) {
+  if (!boolAttr)
+    return false;
+
+  auto type = boolAttr.getType();
+  if (type.isInteger(1)) {
+    auto attr = boolAttr.cast<BoolAttr>();
+    return attr.getValue() == boolVal;
+  }
+  if (auto vecType = type.cast<VectorType>()) {
+    if (vecType.getElementType().isInteger(1))
+      if (auto attr = boolAttr.dyn_cast<SplatElementsAttr>())
+        return attr.getSplatValue().template cast<BoolAttr>().getValue() ==
+               boolVal;
+  }
+  return false;
+}
+
 // Extracts an element from the given `composite` by following the given
 // `indices`. Returns a null Attribute if error happens.
 static Attribute extractCompositeElement(Attribute composite,
                                          ArrayRef<unsigned> indices) {
   // Check that given composite is a constant.
   if (!composite)
     return {};
   // Return composite itself if we reach the end of the index chain.
OpFoldResult spirv::ISubOp::fold(ArrayRef<Attribute> operands) {
   // The resulting value will equal the low-order N bits of the correct result
   // R, where N is the component width and R is computed with enough precision
   // to avoid overflow and underflow.
   return constFoldBinaryOp<IntegerAttr>(operands,
                                         [](APInt a, APInt b) { return a - b; });
 }
 
 //===----------------------------------------------------------------------===//
+// spv.LogicalAnd
+//===----------------------------------------------------------------------===//
+
+OpFoldResult spirv::LogicalAndOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.size() == 2 && "spv.LogicalAnd should take two operands");
+
+  // x && true = x
+  if (isScalarOrSplatBoolAttr(operands.back(), true))
+    return operand1();
+
+  // x && false = false
+  if (isScalarOrSplatBoolAttr(operands.back(), false))
+    return operands.back();
+
+  return Attribute();
+}
+
+//===----------------------------------------------------------------------===//
 // spv.LogicalNot
 //===----------------------------------------------------------------------===//
 
 void spirv::LogicalNotOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<ConvertLogicalNotOfIEqual, ConvertLogicalNotOfINotEqual,
                  ConvertLogicalNotOfLogicalEqual,
                  ConvertLogicalNotOfLogicalNotEqual>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// spv.LogicalOr
+//===----------------------------------------------------------------------===//
+
+OpFoldResult spirv::LogicalOrOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.size() == 2 && "spv.LogicalOr should take two operands");
+
+  // x || true = true
+  if (isScalarOrSplatBoolAttr(operands.back(), true))
+    return operands.back();
+
+  // x || false = x
+  if (isScalarOrSplatBoolAttr(operands.back(), false))
+    return operand1();
+
+  return Attribute();
+}
+
+//===----------------------------------------------------------------------===//
 // spv.selection
 //===----------------------------------------------------------------------===//
 
 namespace {
 // Blocks from the given `spv.selection` operation must satisfy the following
 // layout:
 //
 //       +-----------------------------------------------+

mlir/test/Dialect/SPIRV/canonicalize.mlir

func @const_fold_vector_isub() -> vector<3xi32> {
   // CHECK: spv.constant dense<[45, -40, 99]>
   %0 = spv.ISub %vc1, %vc2 : vector<3xi32>
   return %0: vector<3xi32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
+// spv.LogicalAnd
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @convert_logical_and_true_false_scalar
+// CHECK-SAME: %[[ARG:.+]]: i1
+func @convert_logical_and_true_false_scalar(%arg: i1) -> (i1, i1) {
+  %true = spv.constant true
+  // CHECK: %[[FALSE:.+]] = spv.constant false
+  %false = spv.constant false
+  %0 = spv.LogicalAnd %true, %arg: i1
+  %1 = spv.LogicalAnd %arg, %false: i1
+  // CHECK: return %[[ARG]], %[[FALSE]]
+  return %0, %1: i1, i1
+}
+
+// CHECK-LABEL: @convert_logical_and_true_false_vector
+// CHECK-SAME: %[[ARG:.+]]: vector<3xi1>
+func @convert_logical_and_true_false_vector(%arg: vector<3xi1>) -> (vector<3xi1>, vector<3xi1>) {
+  %true = spv.constant dense<true> : vector<3xi1>
+  // CHECK: %[[FALSE:.+]] = spv.constant dense<false>
+  %false = spv.constant dense<false> : vector<3xi1>
+  %0 = spv.LogicalAnd %true, %arg: vector<3xi1>
+  %1 = spv.LogicalAnd %arg, %false: vector<3xi1>
+  // CHECK: return %[[ARG]], %[[FALSE]]
+  return %0, %1: vector<3xi1>, vector<3xi1>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
 // spv.LogicalNot
 //===----------------------------------------------------------------------===//
 
 func @convert_logical_not_to_not_equal(%arg0: vector<3xi64>, %arg1: vector<3xi64>) -> vector<3xi1> {
   // CHECK: %[[RESULT:.*]] = spv.INotEqual {{%.*}}, {{%.*}} : vector<3xi64>
   // CHECK-NEXT: spv.ReturnValue %[[RESULT]] : vector<3xi1>
   %2 = spv.IEqual %arg0, %arg1 : vector<3xi64>
   %3 = spv.LogicalNot %2 : vector<3xi1>
func @convert_logical_not_to_logical_equal(%arg0: vector<3xi1>, %arg1: vector<3xi1>) -> vector<3xi1> {
   %2 = spv.LogicalNotEqual %arg0, %arg1 : vector<3xi1>
   %3 = spv.LogicalNot %2 : vector<3xi1>
   spv.ReturnValue %3 : vector<3xi1>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
+// spv.LogicalOr
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @convert_logical_or_true_false_scalar
+// CHECK-SAME: %[[ARG:.+]]: i1
+func @convert_logical_or_true_false_scalar(%arg: i1) -> (i1, i1) {
+  // CHECK: %[[TRUE:.+]] = spv.constant true
+  %true = spv.constant true
+  %false = spv.constant false
+  %0 = spv.LogicalOr %true, %arg: i1
+  %1 = spv.LogicalOr %arg, %false: i1
+  // CHECK: return %[[TRUE]], %[[ARG]]
+  return %0, %1: i1, i1
+}
+
+// CHECK-LABEL: @convert_logical_or_true_false_vector
+// CHECK-SAME: %[[ARG:.+]]: vector<3xi1>
+func @convert_logical_or_true_false_vector(%arg: vector<3xi1>) -> (vector<3xi1>, vector<3xi1>) {
+  // CHECK: %[[TRUE:.+]] = spv.constant dense<true>
+  %true = spv.constant dense<true> : vector<3xi1>
+  %false = spv.constant dense<false> : vector<3xi1>
+  %0 = spv.LogicalOr %true, %arg: vector<3xi1>
+  %1 = spv.LogicalOr %arg, %false: vector<3xi1>
+  // CHECK: return %[[TRUE]], %[[ARG]]
+  return %0, %1: vector<3xi1>, vector<3xi1>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
 // spv.selection
 //===----------------------------------------------------------------------===//
 
 func @canonicalize_selection_op_scalar_type(%cond: i1) -> () {
   %0 = spv.constant 0: i32
   // CHECK: %[[TRUE_VALUE:.*]] = spv.constant 1 : i32
   %1 = spv.constant 1: i32
   // CHECK: %[[FALSE_VALUE:.*]] = spv.constant 2 : i32