[Mlir-commits] [mlir] 99fc000 - [MLIR] Expose atomicrmw and/or

[William S. Moses]

Author: William S. Moses
Date: 2021-12-29T00:23:28-05:00
New Revision: 99fc000c87c0ce3a5f698cfa5a67b177ad0cf5f8

URL: https://github.com/llvm/llvm-project/commit/99fc000c87c0ce3a5f698cfa5a67b177ad0cf5f8
DIFF: https://github.com/llvm/llvm-project/commit/99fc000c87c0ce3a5f698cfa5a67b177ad0cf5f8.diff

LOG: [MLIR] Expose atomicrmw and/or

LLVM (dialect and IR) have atomics for and/or. This patch enables atomic_rmw ops in the standard dialect for and/or that lower to these (in addition to the existing atomics such as addi, etc).

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D116345

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/StandardOps/IR/StandardOpsBase.td
    mlir/lib/Analysis/AffineAnalysis.cpp
    mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
    mlir/lib/Dialect/StandardOps/IR/Ops.cpp
    mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/StandardOps/IR/StandardOpsBase.td b/mlir/include/mlir/Dialect/StandardOps/IR/StandardOpsBase.td
index 802a32fce370a..3016a197df0d0 100644
--- a/mlir/include/mlir/Dialect/StandardOps/IR/StandardOpsBase.td
+++ b/mlir/include/mlir/Dialect/StandardOps/IR/StandardOpsBase.td
@@ -26,13 +26,16 @@ def ATOMIC_RMW_KIND_MINS    : I64EnumAttrCase<"mins", 7>;
 def ATOMIC_RMW_KIND_MINU    : I64EnumAttrCase<"minu", 8>;
 def ATOMIC_RMW_KIND_MULF    : I64EnumAttrCase<"mulf", 9>;
 def ATOMIC_RMW_KIND_MULI    : I64EnumAttrCase<"muli", 10>;
+def ATOMIC_RMW_KIND_ORI     : I64EnumAttrCase<"ori", 11>;
+def ATOMIC_RMW_KIND_ANDI    : I64EnumAttrCase<"andi", 12>;
 
 def AtomicRMWKindAttr : I64EnumAttr<
     "AtomicRMWKind", "",
     [ATOMIC_RMW_KIND_ADDF, ATOMIC_RMW_KIND_ADDI, ATOMIC_RMW_KIND_ASSIGN,
      ATOMIC_RMW_KIND_MAXF, ATOMIC_RMW_KIND_MAXS, ATOMIC_RMW_KIND_MAXU,
      ATOMIC_RMW_KIND_MINF, ATOMIC_RMW_KIND_MINS, ATOMIC_RMW_KIND_MINU,
-     ATOMIC_RMW_KIND_MULF, ATOMIC_RMW_KIND_MULI]> {
+     ATOMIC_RMW_KIND_MULF, ATOMIC_RMW_KIND_MULI, ATOMIC_RMW_KIND_ORI,
+     ATOMIC_RMW_KIND_ANDI]> {
   let cppNamespace = "::mlir";
 }
 

diff  --git a/mlir/lib/Analysis/AffineAnalysis.cpp b/mlir/lib/Analysis/AffineAnalysis.cpp
index 873d9b9aa3b49..79a367e337133 100644
--- a/mlir/lib/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Analysis/AffineAnalysis.cpp
@@ -57,6 +57,8 @@ static Value getSupportedReduction(AffineForOp forOp, unsigned pos,
           .Case([](arith::AddFOp) { return AtomicRMWKind::addf; })
           .Case([](arith::MulFOp) { return AtomicRMWKind::mulf; })
           .Case([](arith::AddIOp) { return AtomicRMWKind::addi; })
+          .Case([](arith::AndIOp) { return AtomicRMWKind::andi; })
+          .Case([](arith::OrIOp) { return AtomicRMWKind::ori; })
           .Case([](arith::MulIOp) { return AtomicRMWKind::muli; })
           .Case([](arith::MinFOp) { return AtomicRMWKind::minf; })
           .Case([](arith::MaxFOp) { return AtomicRMWKind::maxf; })

diff  --git a/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp b/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
index f588521ac6ef0..da429dd8af119 100644
--- a/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
+++ b/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
@@ -792,6 +792,10 @@ static Optional<LLVM::AtomicBinOp> matchSimpleAtomicOp(AtomicRMWOp atomicOp) {
     return LLVM::AtomicBinOp::min;
   case AtomicRMWKind::minu:
     return LLVM::AtomicBinOp::umin;
+  case AtomicRMWKind::ori:
+    return LLVM::AtomicBinOp::_or;
+  case AtomicRMWKind::andi:
+    return LLVM::AtomicBinOp::_and;
   default:
     return llvm::None;
   }

diff  --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
index 4c0f69e2ce0a8..02d54472baf50 100644
--- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
@@ -155,6 +155,8 @@ static LogicalResult verify(AtomicRMWOp op) {
   case AtomicRMWKind::mins:
   case AtomicRMWKind::minu:
   case AtomicRMWKind::muli:
+  case AtomicRMWKind::ori:
+  case AtomicRMWKind::andi:
     if (!op.getValue().getType().isa<IntegerType>())
       return op.emitOpError()
              << "with kind '" << stringifyAtomicRMWKind(op.getKind())
@@ -178,7 +180,12 @@ Attribute mlir::getIdentityValueAttr(AtomicRMWKind kind, Type resultType,
   case AtomicRMWKind::addf:
   case AtomicRMWKind::addi:
   case AtomicRMWKind::maxu:
+  case AtomicRMWKind::ori:
     return builder.getZeroAttr(resultType);
+  case AtomicRMWKind::andi:
+    return builder.getIntegerAttr(
+        resultType,
+        APInt::getAllOnes(resultType.cast<IntegerType>().getWidth()));
   case AtomicRMWKind::maxs:
     return builder.getIntegerAttr(
         resultType,
@@ -240,6 +247,10 @@ Value mlir::getReductionOp(AtomicRMWKind op, OpBuilder &builder, Location loc,
     return builder.create<arith::MaxUIOp>(loc, lhs, rhs);
   case AtomicRMWKind::minu:
     return builder.create<arith::MinUIOp>(loc, lhs, rhs);
+  case AtomicRMWKind::ori:
+    return builder.create<arith::OrIOp>(loc, lhs, rhs);
+  case AtomicRMWKind::andi:
+    return builder.create<arith::AndIOp>(loc, lhs, rhs);
   // TODO: Add remaining reduction operations.
   default:
     (void)emitOptionalError(loc, "Reduction operation type not supported");

diff  --git a/mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir b/mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir
index a5c88455124a7..c3282e1903d6f 100644
--- a/mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir
+++ b/mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir
@@ -502,6 +502,10 @@ func @atomic_rmw(%I : memref<10xi32>, %ival : i32, %F : memref<10xf32>, %fval :
   // CHECK: llvm.atomicrmw umin %{{.*}}, %{{.*}} acq_rel
   atomic_rmw addf %fval, %F[%i] : (f32, memref<10xf32>) -> f32
   // CHECK: llvm.atomicrmw fadd %{{.*}}, %{{.*}} acq_rel
+  atomic_rmw ori %ival, %I[%i] : (i32, memref<10xi32>) -> i32
+  // CHECK: llvm.atomicrmw _or %{{.*}}, %{{.*}} acq_rel
+  atomic_rmw andi %ival, %I[%i] : (i32, memref<10xi32>) -> i32
+  // CHECK: llvm.atomicrmw _and %{{.*}}, %{{.*}} acq_rel
   return
 }