[mlir][spirv] Add folds for `BitwiseAnd` and `BitwiseOr`

[kuhar]

Summary

The newly folded patterns get produced by type emulation code, so we
cannot rely on the Arith folders handling these folds.

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

let description = [{
 
     #### Example:
 
     ```mlir
     %2 = spirv.BitwiseAnd %0, %1 : i32
     %2 = spirv.BitwiseAnd %0, %1 : vector<4xi32>
     ```
   }];
+
+  let hasFolder = 1;
 }
 
 // -----
 
 def SPIRV_BitwiseOrOp : SPIRV_BitBinaryOp<"BitwiseOr",
                                       [Commutative, UsableInSpecConstantOp]> {
   let summary = [{
     Result is 1 if either Operand 1 or Operand 2 is 1. Result is 0 if both
let description = [{
 
     #### Example:
 
     ```mlir
     %2 = spirv.BitwiseOr %0, %1 : i32
     %2 = spirv.BitwiseOr %0, %1 : vector<4xi32>
     ```
   }];
+
+  let hasFolder = 1;
 }
 
 // -----
 
 def SPIRV_BitwiseXorOp : SPIRV_BitBinaryOp<"BitwiseXor",
                                        [Commutative, UsableInSpecConstantOp]> {
   let summary = [{
     Result is 1 if exactly one of Operand 1 or Operand 2 is 1. Result is 0

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

 #include "mlir/Interfaces/FunctionImplementation.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringExtras.h"
 #include <cassert>
 #include <numeric>
+#include <optional>
 #include <type_traits>
 
 using namespace mlir;
 using namespace mlir::spirv::AttrNames;
 
 //===----------------------------------------------------------------------===//
 // Common utility functions
 //===----------------------------------------------------------------------===//

 // spirv.ShiftRightLogicalOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult spirv::ShiftRightLogicalOp::verify() {
   return verifyShiftOp(*this);
 }
 
 //===----------------------------------------------------------------------===//
+// spirv.BtiwiseAndOp
+//===----------------------------------------------------------------------===//
+
+static std::optional<APInt> extractIntConstant(Attribute attr) {
+  IntegerAttr intAttr;
+  if (auto splat = dyn_cast_if_present<SplatElementsAttr>(attr))
+    intAttr = dyn_cast<IntegerAttr>(splat.getSplatValue<Attribute>());
+  else
+    intAttr = dyn_cast_if_present<IntegerAttr>(attr);
+
+  if (!intAttr)
+    return std::nullopt;
+
+  return intAttr.getValue();
+}
+
+OpFoldResult
+spirv::BitwiseAndOp::fold(spirv::BitwiseAndOp::FoldAdaptor adaptor) {
+  std::optional<APInt> rhsVal = extractIntConstant(adaptor.getOperand2());
+  if (!rhsVal)
+    return {};
+
+  APInt rhsMask = *rhsVal;
+
+  // x & 0 -> 0
+  if (rhsMask.isZero())
+    return getOperand2();
+
+  // x & <all ones> -> x
+  if (rhsMask.isAllOnes())
+    return getOperand1();
+
+  // (UConvert x : iN to iK) & <mask with N low bits set> -> UConvert x
+  if (auto zext = getOperand1().getDefiningOp<spirv::UConvertOp>()) {
+    int valueBits =
+        getElementTypeOrSelf(zext.getOperand()).getIntOrFloatBitWidth();
+    if (rhsMask.zextOrTrunc(valueBits).isAllOnes())
+      return getOperand1();
+  }
+
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
+// spirv.BtiwiseOrOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult spirv::BitwiseOrOp::fold(spirv::BitwiseOrOp::FoldAdaptor adaptor) {
+  std::optional<APInt> rhsVal = extractIntConstant(adaptor.getOperand2());
+  if (!rhsVal)
+    return {};
+
+  APInt rhsMask = *rhsVal;
+
+  // x | 0 -> x
+  if (rhsMask.isZero())
+    return getOperand1();
+
+  // x | <all ones> -> <all ones>
+  if (rhsMask.isAllOnes())
+    return getOperand2();
+
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
 // spirv.ImageQuerySize
 //===----------------------------------------------------------------------===//
 
 LogicalResult spirv::ImageQuerySizeOp::verify() {
   spirv::ImageType imageType =
       llvm::cast<spirv::ImageType>(getImage().getType());
   Type resultType = getResult().getType();
 

mlir/test/Dialect/SPIRV/IR/bit-ops.mlir

-// RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
+// RUN: mlir-opt --split-input-file --verify-diagnostics %s | FileCheck %s
+// RUN: mlir-opt --split-input-file --verify-diagnostics --canonicalize %s \
+// RUN:  | FileCheck %s --check-prefix=CANON
 
 //===----------------------------------------------------------------------===//
 // spirv.BitCount
 //===----------------------------------------------------------------------===//
 
 func.func @bitcount(%arg: i32) -> i32 {
   // CHECK: spirv.BitCount {{%.*}} : i32
   %0 = spirv.BitCount %arg : i32

 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // spirv.BitwiseOr
 //===----------------------------------------------------------------------===//
 
+// CHECK-LABEL: func @bitwise_or_scalar
 func.func @bitwise_or_scalar(%arg: i32) -> i32 {
   // CHECK: spirv.BitwiseOr
   %0 = spirv.BitwiseOr %arg, %arg : i32
   return %0 : i32
 }
 
+// CHECK-LABEL: func @bitwise_or_vector
 func.func @bitwise_or_vector(%arg: vector<4xi32>) -> vector<4xi32> {
   // CHECK: spirv.BitwiseOr
   %0 = spirv.BitwiseOr %arg, %arg : vector<4xi32>
   return %0 : vector<4xi32>
 }
 
+// CANON-LABEL: func @bitwise_or_zero
+// CANON-SAME:    (%[[ARG:.+]]: i32)
+func.func @bitwise_or_zero(%arg: i32) -> i32 {
+  // CANON: return %[[ARG]]
+  %zero = spirv.Constant 0 : i32
+  %0 = spirv.BitwiseOr %arg, %zero : i32
+  return %0 : i32
+}
+
+// CANON-LABEL: func @bitwise_or_zero_vector
+// CANON-SAME:    (%[[ARG:.+]]: vector<4xi32>)
+func.func @bitwise_or_zero_vector(%arg: vector<4xi32>) -> vector<4xi32> {
+  // CANON: return %[[ARG]]
+  %zero = spirv.Constant dense<0> : vector<4xi32>
+  %0 = spirv.BitwiseOr %arg, %zero : vector<4xi32>
+  return %0 : vector<4xi32>
+}
+
+// CANON-LABEL: func @bitwise_or_all_ones
+func.func @bitwise_or_all_ones(%arg: i8) -> i8 {
+  // CANON: %[[CST:.+]] = spirv.Constant -1
+  // CANON: return %[[CST]]
+  %ones = spirv.Constant 255 : i8
+  %0 = spirv.BitwiseOr %arg, %ones : i8
+  return %0 : i8
+}
+
+// CANON-LABEL: func @bitwise_or_all_ones_vector
+func.func @bitwise_or_all_ones_vector(%arg: vector<3xi8>) -> vector<3xi8> {
+  // CANON: %[[CST:.+]] = spirv.Constant dense<-1>
+  // CANON: return %[[CST]]
+  %ones = spirv.Constant dense<255> : vector<3xi8>
+  %0 = spirv.BitwiseOr %arg, %ones : vector<3xi8>
+  return %0 : vector<3xi8>
+}
+
 // -----
 
 func.func @bitwise_or_float(%arg0: f16, %arg1: f16) -> f16 {
   // expected-error @+1 {{operand #0 must be 8/16/32/64-bit integer or vector of 8/16/32/64-bit integer values of length 2/3/4}}
   %0 = spirv.BitwiseOr %arg0, %arg1 : f16
   return %0 : f16
 }
 

 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // spirv.BitwiseAnd
 //===----------------------------------------------------------------------===//
 
+// CHECK-LABEL: func @bitwise_and_scalar
 func.func @bitwise_and_scalar(%arg: i32) -> i32 {
   // CHECK: spirv.BitwiseAnd
   %0 = spirv.BitwiseAnd %arg, %arg : i32
   return %0 : i32
 }
 
+// CHECK-LABEL: func @bitwise_and_vector
 func.func @bitwise_and_vector(%arg: vector<4xi32>) -> vector<4xi32> {
   // CHECK: spirv.BitwiseAnd
   %0 = spirv.BitwiseAnd %arg, %arg : vector<4xi32>
   return %0 : vector<4xi32>
 }
 
+// CANON-LABEL: func @bitwise_and_zero
+func.func @bitwise_and_zero(%arg: i32) -> i32 {
+  // CANON: %[[CST:.+]] = spirv.Constant 0
+  // CANON: return %[[CST]]
+  %zero = spirv.Constant 0 : i32
+  %0 = spirv.BitwiseAnd %arg, %zero : i32
+  return %0 : i32
+}
+
+// CANON-LABEL: func @bitwise_and_zero_vector
+func.func @bitwise_and_zero_vector(%arg: vector<4xi32>) -> vector<4xi32> {
+  // CANON: %[[CST:.+]] = spirv.Constant dense<0>
+  // CANON: return %[[CST]]
+  %zero = spirv.Constant dense<0> : vector<4xi32>
+  %0 = spirv.BitwiseAnd %arg, %zero : vector<4xi32>
+  return %0 : vector<4xi32>
+}
+
+// CANON-LABEL: func @bitwise_and_all_ones
+// CANON-SAME:    (%[[ARG:.+]]: i8)
+func.func @bitwise_and_all_ones(%arg: i8) -> i8 {
+  // CANON: return %[[ARG]]
+  %ones = spirv.Constant 255 : i8
+  %0 = spirv.BitwiseAnd %arg, %ones : i8
+  return %0 : i8
+}
+
+// CANON-LABEL: func @bitwise_and_all_ones_vector
+// CANON-SAME:    (%[[ARG:.+]]: vector<3xi8>)
+func.func @bitwise_and_all_ones_vector(%arg: vector<3xi8>) -> vector<3xi8> {
+  // CANON: return %[[ARG]]
+  %ones = spirv.Constant dense<255> : vector<3xi8>
+  %0 = spirv.BitwiseAnd %arg, %ones : vector<3xi8>
+  return %0 : vector<3xi8>
+}
+
+// CANON-LABEL: func @bitwise_and_zext_1
+// CANON-SAME:    (%[[ARG:.+]]: i8)
+func.func @bitwise_and_zext_1(%arg: i8) -> i32 {
+  // CANON: %[[ZEXT:.+]] = spirv.UConvert %[[ARG]]
+  // CANON: return %[[ZEXT]]
+  %zext = spirv.UConvert %arg : i8 to i32
+  %ones = spirv.Constant 255 : i32
+  %0 = spirv.BitwiseAnd %zext, %ones : i32
+  return %0 : i32
+}
+
+// CANON-LABEL: func @bitwise_and_zext_2
+// CANON-SAME:    (%[[ARG:.+]]: i8)
+func.func @bitwise_and_zext_2(%arg: i8) -> i32 {
+  // CANON: %[[ZEXT:.+]] = spirv.UConvert %[[ARG]]
+  // CANON: return %[[ZEXT]]
+  %zext = spirv.UConvert %arg : i8 to i32
+  %ones = spirv.Constant 0x12345ff : i32
+  %0 = spirv.BitwiseAnd %zext, %ones : i32
+  return %0 : i32
+}
+
+// CANON-LABEL: func @bitwise_and_zext_3
+// CANON-SAME:    (%[[ARG:.+]]: i8)
+func.func @bitwise_and_zext_3(%arg: i8) -> i32 {
+  // CANON: %[[ZEXT:.+]] = spirv.UConvert %[[ARG]]
+  // CANON: %[[AND:.+]]  = spirv.BitwiseAnd %[[ZEXT]]
+  // CANON: return %[[AND]]
+  %zext = spirv.UConvert %arg : i8 to i32
+  %ones = spirv.Constant 254 : i32
+  %0 = spirv.BitwiseAnd %zext, %ones : i32
+  return %0 : i32
+}
+
+// CANON-LABEL: func @bitwise_and_zext_vector
+// CANON-SAME:    (%[[ARG:.+]]: vector<2xi8>)
+func.func @bitwise_and_zext_vector(%arg: vector<2xi8>) -> vector<2xi32> {
+  // CANON: %[[ZEXT:.+]] = spirv.UConvert %[[ARG]]
+  // CANON: return %[[ZEXT]]
+  %zext = spirv.UConvert %arg : vector<2xi8> to vector<2xi32>
+  %ones = spirv.Constant dense<255> : vector<2xi32>
+  %0 = spirv.BitwiseAnd %zext, %ones : vector<2xi32>
+  return %0 : vector<2xi32>
+}
+
 // -----
 
 func.func @bitwise_and_float(%arg0: f16, %arg1: f16) -> f16 {
   // expected-error @+1 {{operand #0 must be 8/16/32/64-bit integer or vector of 8/16/32/64-bit integer values of length 2/3/4}}
   %0 = spirv.BitwiseAnd %arg0, %arg1 : f16
   return %0 : f16
 }