[MLIR] Expose atomicrmw and/or

[wsmoses]

Summary

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).

[wsmoses]

Remove excess code

[mehdi_amini]

LGTM.
I wonder why isn't this operation didn't move in the arith dialect as well?

mlir/lib/Analysis/AffineAnalysis.cpp

1

Nit: you should install clang-format on your machine

[wsmoses]

Yeah I concur that this (and the other atomics) would be wise to move.

mlir/include/mlir/Dialect/StandardOps/IR/StandardOpsBase.td

 def ATOMIC_RMW_KIND_MAXF    : I64EnumAttrCase<"maxf", 3>;
 def ATOMIC_RMW_KIND_MAXS    : I64EnumAttrCase<"maxs", 4>;
 def ATOMIC_RMW_KIND_MAXU    : I64EnumAttrCase<"maxu", 5>;
 def ATOMIC_RMW_KIND_MINF    : I64EnumAttrCase<"minf", 6>;
 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";
 }
 
 #endif // STANDARD_OPS_BASE

mlir/lib/Analysis/AffineAnalysis.cpp

 //===- AffineAnalysis.cpp - Affine structures analysis routines -----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements miscellaneous analysis routines for affine structures
if (combinerOps.size() > 1)
     return nullptr;
 
   Operation *combinerOp = combinerOps.back();
   Optional<AtomicRMWKind> maybeKind =
       TypeSwitch<Operation *, Optional<AtomicRMWKind>>(combinerOp)
           .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; })
           .Case([](arith::MinSIOp) { return AtomicRMWKind::mins; })
           .Case([](arith::MaxSIOp) { return AtomicRMWKind::maxs; })
           .Case([](arith::MinUIOp) { return AtomicRMWKind::minu; })
           .Case([](arith::MaxUIOp) { return AtomicRMWKind::maxu; })
           .Default([](Operation *) -> Optional<AtomicRMWKind> {

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

static Optional<LLVM::AtomicBinOp> matchSimpleAtomicOp(AtomicRMWOp atomicOp) {
   case AtomicRMWKind::maxs:
     return LLVM::AtomicBinOp::max;
   case AtomicRMWKind::maxu:
     return LLVM::AtomicBinOp::umax;
   case AtomicRMWKind::mins:
     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;
   }
   llvm_unreachable("Invalid AtomicRMWKind");
 }
 
 namespace {
 
 struct AtomicRMWOpLowering : public LoadStoreOpLowering<AtomicRMWOp> {
   using Base::Base;
 
   LogicalResult
   matchAndRewrite(AtomicRMWOp atomicOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     if (failed(match(atomicOp)))
       return failure();
     auto maybeKind = matchSimpleAtomicOp(atomicOp);
-    if (!maybeKind)
+    if (!maybeKind) {
       return failure();
+    }
     auto resultType = adaptor.getValue().getType();
     auto memRefType = atomicOp.getMemRefType();
     auto dataPtr =
         getStridedElementPtr(atomicOp.getLoc(), memRefType, adaptor.getMemref(),
                              adaptor.getIndices(), rewriter);
     rewriter.replaceOpWithNewOp<LLVM::AtomicRMWOp>(
         atomicOp, resultType, *maybeKind, dataPtr, adaptor.getValue(),
         LLVM::AtomicOrdering::acq_rel);

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

if (!op.getValue().getType().isa<FloatType>())
              << "' expects a floating-point type";
     break;
   case AtomicRMWKind::addi:
   case AtomicRMWKind::maxs:
   case AtomicRMWKind::maxu:
   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())
              << "' expects an integer type";
     break;
   default:
     break;
   }
   return success();
 }
 
 /// Returns the identity value attribute associated with an AtomicRMWKind op.
 Attribute mlir::getIdentityValueAttr(AtomicRMWKind kind, Type resultType,
                                      OpBuilder &builder, Location loc) {
   switch (kind) {
   case AtomicRMWKind::maxf:
     return builder.getFloatAttr(
         resultType,
         APFloat::getInf(resultType.cast<FloatType>().getFloatSemantics(),
                         /*Negative=*/true));
   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,
         APInt::getSignedMinValue(resultType.cast<IntegerType>().getWidth()));
   case AtomicRMWKind::minf:
     return builder.getFloatAttr(
         resultType,
         APFloat::getInf(resultType.cast<FloatType>().getFloatSemantics(),
Value mlir::getReductionOp(AtomicRMWKind op, OpBuilder &builder, Location loc,
   case AtomicRMWKind::maxs:
     return builder.create<arith::MaxSIOp>(loc, lhs, rhs);
   case AtomicRMWKind::mins:
     return builder.create<arith::MinSIOp>(loc, lhs, rhs);
   case AtomicRMWKind::maxu:
     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");
     break;
   }
   return nullptr;
 }
 

mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir

func @atomic_rmw(%I : memref<10xi32>, %ival : i32, %F : memref<10xf32>, %fval : f32, %i : index) {
   atomic_rmw mins %ival, %I[%i] : (i32, memref<10xi32>) -> i32
   // CHECK: llvm.atomicrmw min %{{.*}}, %{{.*}} acq_rel
   atomic_rmw maxu %ival, %I[%i] : (i32, memref<10xi32>) -> i32
   // CHECK: llvm.atomicrmw umax %{{.*}}, %{{.*}} acq_rel
   atomic_rmw minu %ival, %I[%i] : (i32, memref<10xi32>) -> i32
   // 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
 }
 
 // -----
 
 // CHECK-LABEL: func @generic_atomic_rmw
 func @generic_atomic_rmw(%I : memref<10xi32>, %i : index) -> i32 {
   %x = generic_atomic_rmw %I[%i] : memref<10xi32> {