From 4f9c9295a6057d752c1fe0220d797ee94b0a18cf Mon Sep 17 00:00:00 2001
From: Luca Boasso <luca@modular.com>
Date: Wed, 23 Nov 2022 13:26:02 -0600
Subject: [PATCH] [mlir][index] Add and, or, and xor ops

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

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D138590
---
 mlir/include/mlir/Dialect/Index/IR/IndexOps.td     | 57 ++++++++++++++++++++++
 mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp    |  6 +++
 mlir/lib/Dialect/Index/IR/IndexOps.cpp             | 27 ++++++++++
 .../test/Conversion/IndexToLLVM/index-to-llvm.mlir |  8 ++-
 mlir/test/Dialect/Index/index-canonicalize.mlir    | 30 ++++++++++++
 mlir/test/Dialect/Index/index-ops.mlir             |  6 +++
 6 files changed, 133 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Index/IR/IndexOps.td b/mlir/include/mlir/Dialect/Index/IR/IndexOps.td
index 29f4c1e..233b752 100644
--- a/mlir/include/mlir/Dialect/Index/IR/IndexOps.td
+++ b/mlir/include/mlir/Dialect/Index/IR/IndexOps.td
@@ -344,6 +344,63 @@ def Index_ShrUOp : IndexBinaryOp<"shru"> {
 }
 
 //===----------------------------------------------------------------------===//
+// 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
 //===----------------------------------------------------------------------===//
 
diff --git a/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp b/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp
index 4461d51..9fa2e53 100644
--- a/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp
+++ b/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp
@@ -273,6 +273,9 @@ 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>;
 
@@ -298,6 +301,9 @@ void index::populateIndexToLLVMConversionPatterns(
       ConvertIndexShl,
       ConvertIndexShrS,
       ConvertIndexShrU,
+      ConvertIndexAnd,
+      ConvertIndexOr,
+      ConvertIndexXor,
       ConvertIndexCeilDivS,
       ConvertIndexCeilDivU,
       ConvertIndexFloorDivS,
diff --git a/mlir/lib/Dialect/Index/IR/IndexOps.cpp b/mlir/lib/Dialect/Index/IR/IndexOps.cpp
index 2178a75..00e31bc 100644
--- a/mlir/lib/Dialect/Index/IR/IndexOps.cpp
+++ b/mlir/lib/Dialect/Index/IR/IndexOps.cpp
@@ -330,6 +330,33 @@ OpFoldResult ShrUOp::fold(ArrayRef<Attribute> operands) {
 }
 
 //===----------------------------------------------------------------------===//
+// 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
 //===----------------------------------------------------------------------===//
 
diff --git a/mlir/test/Conversion/IndexToLLVM/index-to-llvm.mlir b/mlir/test/Conversion/IndexToLLVM/index-to-llvm.mlir
index c6b2273..44aea80 100644
--- a/mlir/test/Conversion/IndexToLLVM/index-to-llvm.mlir
+++ b/mlir/test/Conversion/IndexToLLVM/index-to-llvm.mlir
@@ -28,8 +28,14 @@ func.func @trivial_ops(%a: index, %b: index) {
   %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
 }
 
diff --git a/mlir/test/Dialect/Index/index-canonicalize.mlir b/mlir/test/Dialect/Index/index-canonicalize.mlir
index 288593f..d525ecd 100644
--- a/mlir/test/Dialect/Index/index-canonicalize.mlir
+++ b/mlir/test/Dialect/Index/index-canonicalize.mlir
@@ -384,6 +384,36 @@ func.func @shru_edge() -> index {
   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
diff --git a/mlir/test/Dialect/Index/index-ops.mlir b/mlir/test/Dialect/Index/index-ops.mlir
index d1a4097..5fa0498 100644
--- a/mlir/test/Dialect/Index/index-ops.mlir
+++ b/mlir/test/Dialect/Index/index-ops.mlir
@@ -33,6 +33,12 @@ func.func @binary_ops(%a: index, %b: index) {
   %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
 }
 
-- 
2.7.4