[mlir][index] Add and, or, and xor ops

[lboasso]

Summary

This patch adds the and, or, and xor bitwise operations to
the index dialects with folders and LLVM lowerings.

[lboasso]

I uploaded the wrong patch before, this is the correct one.

[rriddle]

mlir/include/mlir/Dialect/Index/IR/IndexOps.td

354

The formatting here looks off.

358–361

373

Same as AndOp, and throughout this commit.

[lboasso]

format

[lboasso]

Sorry my editor tried to format IndexOps.td. It should be better now.

mlir/include/mlir/Dialect/Index/IR/IndexOps.td

let description = [{
     ```mlir
     // c = a >> b
     %c = index.shru %a, %b
     ```
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// AndOp
+//===----------------------------------------------------------------------===//
+
+def Index_AndOp : IndexBinaryOp<"and"> {
+  let summary = "index bitwise and";
+  let description = [{
+    The `index.and` operation takes two index values and computes their bitwise
+    and.
+
+    Example:
+
+    ```mlir
+    // c = a & b
+    %c = index.and %a, %b
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// OrOp
+//===----------------------------------------------------------------------===//
+
+def Index_OrOp : IndexBinaryOp<"or"> {
+  let summary = "index bitwise or";
+  let description = [{
+    The `index.or` operation takes two index values and computes their bitwise
+    or.
+
+    Example:
+
+    ```mlir
+    // c = a | b
+    %c = index.or %a, %b
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// XorOp
+//===----------------------------------------------------------------------===//
+
+def Index_XOrOp : IndexBinaryOp<"xor"> {
+  let summary = "index bitwise xor";
+  let description = [{
+    The `index.xor` operation takes two index values and computes their bitwise
+    xor.
+
+    Example:
+
+    ```mlir
+    // c = a ^ b
+    %c = index.xor %a, %b
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // CastSOp
 //===----------------------------------------------------------------------===//
 
 def Index_CastSOp : IndexOp<"casts",
                            [DeclareOpInterfaceMethods<CastOpInterface>]> {
   let summary = "index signed cast";
   let description = [{
     The `index.casts` operation enables conversions between values of index type

mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp

using ConvertIndexMaxS =
     mlir::OneToOneConvertToLLVMPattern<MaxSOp, LLVM::SMaxOp>;
 using ConvertIndexMaxU =
     mlir::OneToOneConvertToLLVMPattern<MaxUOp, LLVM::UMaxOp>;
 using ConvertIndexShl = mlir::OneToOneConvertToLLVMPattern<ShlOp, LLVM::ShlOp>;
 using ConvertIndexShrS =
     mlir::OneToOneConvertToLLVMPattern<ShrSOp, LLVM::AShrOp>;
 using ConvertIndexShrU =
     mlir::OneToOneConvertToLLVMPattern<ShrUOp, LLVM::LShrOp>;
+using ConvertIndexAnd = mlir::OneToOneConvertToLLVMPattern<AndOp, LLVM::AndOp>;
+using ConvertIndexOr = mlir::OneToOneConvertToLLVMPattern<OrOp, LLVM::OrOp>;
+using ConvertIndexXor = mlir::OneToOneConvertToLLVMPattern<XOrOp, LLVM::XOrOp>;
 using ConvertIndexBoolConstant =
     mlir::OneToOneConvertToLLVMPattern<BoolConstantOp, LLVM::ConstantOp>;
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Pattern Population
 //===----------------------------------------------------------------------===//
patterns.insert<
       ConvertIndexDivU,
       ConvertIndexRemS,
       ConvertIndexRemU,
       ConvertIndexMaxS,
       ConvertIndexMaxU,
       ConvertIndexShl,
       ConvertIndexShrS,
       ConvertIndexShrU,
+      ConvertIndexAnd,
+      ConvertIndexOr,
+      ConvertIndexXor,
       ConvertIndexCeilDivS,
       ConvertIndexCeilDivU,
       ConvertIndexFloorDivS,
       ConvertIndexCastS,
       ConvertIndexCastU,
       ConvertIndexCmp,
       ConvertIndexSizeOf,
       ConvertIndexConstant,

mlir/lib/Dialect/Index/IR/IndexOps.cpp

return foldBinaryOpChecked(
         // Don't fold if RHS is greater than or equal to 32.
         if (rhs.uge(32))
           return {};
         return lhs.lshr(rhs);
       });
 }
 
 //===----------------------------------------------------------------------===//
+// AndOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult AndOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpUnchecked(
+      operands, [](const APInt &lhs, const APInt &rhs) { return lhs & rhs; });
+}
+
+//===----------------------------------------------------------------------===//
+// OrOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult OrOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpUnchecked(
+      operands, [](const APInt &lhs, const APInt &rhs) { return lhs | rhs; });
+}
+
+//===----------------------------------------------------------------------===//
+// XOrOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult XOrOp::fold(ArrayRef<Attribute> operands) {
+  return foldBinaryOpUnchecked(
+      operands, [](const APInt &lhs, const APInt &rhs) { return lhs ^ rhs; });
+}
+
+//===----------------------------------------------------------------------===//
 // CastSOp
 //===----------------------------------------------------------------------===//
 
 bool CastSOp::areCastCompatible(TypeRange lhsTypes, TypeRange rhsTypes) {
   return lhsTypes.front().isa<IndexType>() != rhsTypes.front().isa<IndexType>();
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Conversion/IndexToLLVM/index-to-llvm.mlir

func.func @trivial_ops(%a: index, %b: index) {
   // CHECK: llvm.intr.umax
   %8 = index.maxu %a, %b
   // CHECK: llvm.shl
   %9 = index.shl %a, %b
   // CHECK: llvm.ashr
   %10 = index.shrs %a, %b
   // CHECK: llvm.lshr
   %11 = index.shru %a, %b
+  // CHECK: llvm.add
+  %12 = index.add %a, %b
+  // CHECK: llvm.or
+  %13 = index.or %a, %b
+  // CHECK: llvm.xor
+  %14 = index.xor %a, %b
   // CHECK: llvm.mlir.constant(true
-  %12 = index.bool.constant true
+  %15 = index.bool.constant true
   return
 }
 
 // CHECK-LABEL: @ceildivs
 // CHECK-SAME: %[[NI:.*]]: index, %[[MI:.*]]: index
 func.func @ceildivs(%n: index, %m: index) -> index {
   // CHECK: %[[N:.*]] = builtin.unrealized_conversion_cast %[[NI]]
   // CHECK: %[[M:.*]] = builtin.unrealized_conversion_cast %[[MI]]

mlir/test/Dialect/Index/index-canonicalize.mlir

func.func @shru_edge() -> index {
   %lhs = index.constant 0x10000000000
   %rhs = index.constant 3
   // CHECK: %[[A:.*]] = index.constant 137{{[0-9]+}}
   %0 = index.shru %lhs, %rhs
   // CHECK: return %[[A]]
   return %0 : index
 }
 
+// CHECK-LABEL: @and
+func.func @and() -> index {
+  %lhs = index.constant 5
+  %rhs = index.constant 1
+  // CHECK: %[[A:.*]] = index.constant 1
+  %0 = index.and %lhs, %rhs
+  // CHECK: return %[[A]]
+  return %0 : index
+}
+
+// CHECK-LABEL: @or
+func.func @or() -> index {
+  %lhs = index.constant 5
+  %rhs = index.constant 2
+  // CHECK: %[[A:.*]] = index.constant 7
+  %0 = index.or %lhs, %rhs
+  // CHECK: return %[[A]]
+  return %0 : index
+}
+
+// CHECK-LABEL: @xor
+func.func @xor() -> index {
+  %lhs = index.constant 5
+  %rhs = index.constant 1
+  // CHECK: %[[A:.*]] = index.constant 4
+  %0 = index.xor %lhs, %rhs
+  // CHECK: return %[[A]]
+  return %0 : index
+}
+
 // CHECK-LABEL: @cmp
 func.func @cmp() -> (i1, i1, i1, i1) {
   %a = index.constant 0
   %b = index.constant -1
   %c = index.constant -2
   %d = index.constant 4
 
   %0 = index.cmp slt(%a, %b)

mlir/test/Dialect/Index/index-ops.mlir

func.func @binary_ops(%a: index, %b: index) {
   // CHECK-NEXT: index.maxu %[[A]], %[[B]]
   %11 = index.maxu %a, %b
   // CHECK-NEXT: index.shl %[[A]], %[[B]]
   %12 = index.shl %a, %b
   // CHECK-NEXT: index.shrs %[[A]], %[[B]]
   %13 = index.shrs %a, %b
   // CHECK-NEXT: index.shru %[[A]], %[[B]]
   %14 = index.shru %a, %b
+  // CHECK-NEXT: index.and %[[A]], %[[B]]
+  %15 = index.and %a, %b
+  // CHECK-NEXT: index.or %[[A]], %[[B]]
+  %16 = index.or %a, %b
+  // CHECK-NEXT: index.xor %[[A]], %[[B]]
+  %17 = index.xor %a, %b
   return
 }
 
 // CHECK-LABEL: @cmp_op
 // CHECK-SAME: %[[A:.*]]: index, %[[B:.*]]: index
 func.func @cmp_op(%a: index, %b: index) {
   // CHECK-NEXT: index.cmp eq(%[[A]], %[[B]])
   %0 = index.cmp eq(%a, %b)