[MLIR] Add conversion from AtomicRMWOp -> GenericAtomicRMWOp.

[pifon2a]

Summary

Adding this pattern reduces code duplication. There is no need to have a custom implementation for lowering to llvm.cmpxchg.

[herhut]

Thanks for cleaning this up!

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2646

I would drop this. The pattern already checks for two specific operations. You can make it an assert if you want to ensure that the definition of matchSimpleAtomicOp and this one stay in sync.

Also, this should be in ConvertStandardToStandard (which is possible without the dependency on matchSimpleAtomicOp).

2652

This pattern does not have a test.

Also, maybe call it AtomicMinfMaxfTo... as that is what it is doing.

2851

Expose this from a the ConvertStandardToStandard header file with a descriptive name.

[ftynse]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2851

Yes, I was thinking that this pattern is not at all concerned with LLVM and should not live in "StandardToLLVM". However, the decision when to apply this pattern is driven by what LLVM supports.

Also, ConvertStandardToStandard feels weird to me. lib/Conversion was originally intended to solve the problem of the A->B conversion depending both on A dialect and B dialect so making little sense in either Dialect/A or Dialect/B... I would rather have it in lib/Dialect/StandardOps/Transforms/ExpandAtomic and call it from the LLVM conversion if necessary. If we ever have non-LLVM users or want to use the expansion for other cases than we currently do, it will be easier to extend.

[herhut]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2851

Thinking of it, we do the same in the GPU dialect where we have the AllReduceOpLowering that also lives in Transforms. So +1.

[pifon2a]

Addressed the comments.

[pifon2a]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2652

It was tested through lowering to cmpxchg. Now there is a separate test for this very pattern.

2851

Should we also clean-up mlir/lib/Conversion/StandardToStandard/StandardToStandard.cpp and move it to StandardOps/Transforms?

[pifon2a]

updated the comment.

[herhut]

Nice! Thanks.

mlir/docs/Passes.md

 
 ## `quant` Dialect Passes
 
 [include "QuantPasses.md"]
 
 ## `spv` Dialect Passes
 
 [include "SPIRVPasses.md"]
+
+## `standard` Dialect Passes
+
+[include "StandardPasses.md"]

mlir/include/mlir/Dialect/StandardOps/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(Transforms)

mlir/include/mlir/Dialect/StandardOps/Transforms/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls)
+add_public_tablegen_target(MLIRStandardTransformsIncGen)
+
+add_mlir_doc(Passes -gen-pass-doc StandardPasses ./)

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.h

+
+//===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===//
+//
+// 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 header file defines prototypes that expose pass constructors in the loop
+// transformation library.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_
+#define MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_
+
+#include <memory>
+
+namespace mlir {
+
+class Pass;
+
+/// Creates an instance of the ExpandAtomic pass.
+std::unique_ptr<Pass> createExpandAtomicPass();
+
+} // end namespace mlir
+
+#endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.td

+//===-- Passes.td - StandardOps pass definition file -------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES
+#define MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def ExpandAtomic : FunctionPass<"expand-atomic"> {
+  let summary = "Expands AtomicRMWOp into GenericAtomicRMWOp.";
+  let constructor = "mlir::createExpandAtomicPass()";
+}
+
+#endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES

mlir/include/mlir/InitAllPasses.h

 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/LLVMIR/Transforms/LegalizeForExport.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/LoopOps/Passes.h"
 #include "mlir/Dialect/Quant/Passes.h"
 #include "mlir/Dialect/SPIRV/Passes.h"
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/Transforms/LocationSnapshot.h"
 #include "mlir/Transforms/Passes.h"
 #include "mlir/Transforms/ViewOpGraph.h"
 #include "mlir/Transforms/ViewRegionGraph.h"
 
 #include <cstdlib>
 
 namespace mlir {
#include "mlir/Dialect/LoopOps/Passes.h.inc"
 
   // Quant
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Quant/Passes.h.inc"
 
   // SPIR-V
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/SPIRV/Passes.h.inc"
+
+  // Standard
+#define GEN_PASS_REGISTRATION
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h.inc"
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLPASSES_H_

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     rewriter.replaceOpWithNewOp<LLVM::AtomicRMWOp>(
         op, resultType, *maybeKind, dataPtr, adaptor.value(),
         LLVM::AtomicOrdering::acq_rel);
     return success();
   }
 };
 
 /// Wrap a llvm.cmpxchg operation in a while loop so that the operation can be
 /// retried until it succeeds in atomically storing a new value into memory.
 ///
 ///      +---------------------------------+
 ///      |   <code before the AtomicRMWOp> |
 ///      |   <compute initial %loaded>     |
 ///      |   br loop(%loaded)              |
 ///      +---------------------------------+
 ///             |
 ///  -------|   |
 ///  |      v   v
 ///  |   +--------------------------------+
 ///  |   | loop(%loaded):                 |
 ///  |   |   <body contents>              |
 ///  |   |   %pair = cmpxchg              |
 ///  |   |   %ok = %pair[0]               |
 ///  |   |   %new = %pair[1]              |
 ///  |   |   cond_br %ok, end, loop(%new) |
 ///  |   +--------------------------------+
 ///  |          |        |
 ///  |-----------        |
 ///                      v
 ///      +--------------------------------+
 ///      | end:                           |
 ///      |   <code after the AtomicRMWOp> |
 ///      +--------------------------------+
 ///
-struct AtomicCmpXchgOpLowering : public LoadStoreOpLowering<AtomicRMWOp> {
-  using Base::Base;
-
-  LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const override {
-    auto atomicOp = cast<AtomicRMWOp>(op);
-    auto maybeKind = matchSimpleAtomicOp(atomicOp);
-    if (maybeKind)
-      return failure();
-
-    LLVM::FCmpPredicate predicate;
-    switch (atomicOp.kind()) {
-    case AtomicRMWKind::maxf:
-      predicate = LLVM::FCmpPredicate::ogt;
-      break;
-    case AtomicRMWKind::minf:
-      predicate = LLVM::FCmpPredicate::olt;
-      break;
-    default:
-      return failure();
-    }
-
-    OperandAdaptor<AtomicRMWOp> adaptor(operands);
-    auto loc = op->getLoc();
-    auto valueType = adaptor.value().getType().cast<LLVM::LLVMType>();
-
-    // Split the block into initial, loop, and ending parts.
-    auto *initBlock = rewriter.getInsertionBlock();
-    auto initPosition = rewriter.getInsertionPoint();
-    auto *loopBlock = rewriter.splitBlock(initBlock, initPosition);
-    auto loopArgument = loopBlock->addArgument(valueType);
-    auto loopPosition = rewriter.getInsertionPoint();
-    auto *endBlock = rewriter.splitBlock(loopBlock, loopPosition);
-
-    // Compute the loaded value and branch to the loop block.
-    rewriter.setInsertionPointToEnd(initBlock);
-    auto memRefType = atomicOp.getMemRefType();
-    auto dataPtr = getDataPtr(loc, memRefType, adaptor.memref(),
-                              adaptor.indices(), rewriter, getModule());
-    Value init = rewriter.create<LLVM::LoadOp>(loc, dataPtr);
-    rewriter.create<LLVM::BrOp>(loc, init, loopBlock);
-
-    // Prepare the body of the loop block.
-    rewriter.setInsertionPointToStart(loopBlock);
-    auto predicateI64 =
-        rewriter.getI64IntegerAttr(static_cast<int64_t>(predicate));
-    auto boolType = LLVM::LLVMType::getInt1Ty(&getDialect());
-    auto lhs = loopArgument;
-    auto rhs = adaptor.value();
-    auto cmp =
-        rewriter.create<LLVM::FCmpOp>(loc, boolType, predicateI64, lhs, rhs);
-    auto select = rewriter.create<LLVM::SelectOp>(loc, cmp, lhs, rhs);
-
-    // Prepare the epilog of the loop block.
-    rewriter.setInsertionPointToEnd(loopBlock);
-    // Append the cmpxchg op to the end of the loop block.
-    auto successOrdering = LLVM::AtomicOrdering::acq_rel;
-    auto failureOrdering = LLVM::AtomicOrdering::monotonic;
-    auto pairType = LLVM::LLVMType::getStructTy(valueType, boolType);
-    auto cmpxchg = rewriter.create<LLVM::AtomicCmpXchgOp>(
-        loc, pairType, dataPtr, loopArgument, select, successOrdering,
-        failureOrdering);
-    // Extract the %new_loaded and %ok values from the pair.
-    Value newLoaded = rewriter.create<LLVM::ExtractValueOp>(
-        loc, valueType, cmpxchg, rewriter.getI64ArrayAttr({0}));
-    Value ok = rewriter.create<LLVM::ExtractValueOp>(
-        loc, boolType, cmpxchg, rewriter.getI64ArrayAttr({1}));
-
-    // Conditionally branch to the end or back to the loop depending on %ok.
-    rewriter.create<LLVM::CondBrOp>(loc, ok, endBlock, ArrayRef<Value>(),
-                                    loopBlock, newLoaded);
-
-    // The 'result' of the atomic_rmw op is the newly loaded value.
-    rewriter.replaceOp(op, {newLoaded});
-
-    return success();
-  }
-};
-
-/// Wrap a llvm.cmpxchg operation in a while loop so that the operation can be
-/// retried until it succeeds in atomically storing a new value into memory.
-///
-///      +---------------------------------+
-///      |   <code before the AtomicRMWOp> |
-///      |   <compute initial %loaded>     |
-///      |   br loop(%loaded)              |
-///      +---------------------------------+
-///             |
-///  -------|   |
-///  |      v   v
-///  |   +--------------------------------+
-///  |   | loop(%loaded):                 |
-///  |   |   <body contents>              |
-///  |   |   %pair = cmpxchg              |
-///  |   |   %ok = %pair[0]               |
-///  |   |   %new = %pair[1]              |
-///  |   |   cond_br %ok, end, loop(%new) |
-///  |   +--------------------------------+
-///  |          |        |
-///  |-----------        |
-///                      v
-///      +--------------------------------+
-///      | end:                           |
-///      |   <code after the AtomicRMWOp> |
-///      +--------------------------------+
-///
 struct GenericAtomicRMWOpLowering
     : public LoadStoreOpLowering<GenericAtomicRMWOp> {
   using Base::Base;
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto atomicOp = cast<GenericAtomicRMWOp>(op);
void mlir::populateStdToLLVMNonMemoryConversionPatterns(
   // FIXME: this should be tablegen'ed
   // clang-format off
   patterns.insert<
       AbsFOpLowering,
       AddFOpLowering,
       AddIOpLowering,
       AllocaOpLowering,
       AndOpLowering,
-      AtomicCmpXchgOpLowering,
       AtomicRMWOpLowering,
       BranchOpLowering,
       CallIndirectOpLowering,
       CallOpLowering,
       CeilFOpLowering,
       CmpFOpLowering,
       CmpIOpLowering,
       CondBranchOpLowering,
patterns.insert<
       ViewOpLowering>(converter);
   patterns.insert<
       AllocOpLowering
       >(converter, useAlignedAlloc);
   // clang-format on
 }
 
 void mlir::populateStdToLLVMDefaultFuncOpConversionPattern(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
     bool emitCWrappers) {
   patterns.insert<FuncOpConversion>(converter, emitCWrappers);
 }
 
 void mlir::populateStdToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
     bool emitCWrappers, bool useAlignedAlloc) {
   populateStdToLLVMDefaultFuncOpConversionPattern(converter, patterns,

mlir/lib/Dialect/StandardOps/CMakeLists.txt

target_link_libraries(MLIRStandardOps
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIREDSC
   MLIRIR
   MLIRSideEffects
   MLIRViewLikeInterface
   LLVMSupport
   )
+
+add_subdirectory(Transforms)

mlir/lib/Dialect/StandardOps/Transforms/CMakeLists.txt

+add_mlir_dialect_library(MLIRStandardOpsTransforms
+  ExpandAtomic.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps/Transforms
+
+  DEPENDS
+  MLIRStandardTransformsIncGen
+  )
+target_link_libraries(MLIRStandardOpsTransforms
+  PUBLIC
+  MLIRIR
+  MLIRPass
+  MLIRStandardOps
+  MLIRSupport
+  LLVMSupport
+  )

mlir/lib/Dialect/StandardOps/Transforms/ExpandAtomic.cpp

+//===- ExpandAtomic.cpp - Code to perform loop fusion ---------------------===//
+//
+// 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 expansion of AtomicRMWOp into GenericAtomicRMWOp.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+
+/// Converts `atomic_rmw` that cannot be lowered to a simple atomic op with
+/// AtomicRMWOpLowering pattern, e.g. with "minf" or "maxf" attributes, to
+/// `generic_atomic_rmw` with the expanded code.
+///
+/// %x = atomic_rmw "maxf" %fval, %F[%i] : (f32, memref<10xf32>) -> f32
+///
+/// will be lowered to
+///
+/// %x = std.generic_atomic_rmw %F[%i] : memref<10xf32> {
+/// ^bb0(%current: f32):
+///   %cmp = cmpf "ogt", %current, %fval : f32
+///   %new_value = select %cmp, %current, %fval : f32
+///   atomic_yield %new_value : f32
+/// }
+struct AtomicRMWOpConverter : public OpRewritePattern<AtomicRMWOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AtomicRMWOp op,
+                                PatternRewriter &rewriter) const final {
+    CmpFPredicate predicate;
+    switch (op.kind()) {
+    case AtomicRMWKind::maxf:
+      predicate = CmpFPredicate::OGT;
+      break;
+    case AtomicRMWKind::minf:
+      predicate = CmpFPredicate::OLT;
+      break;
+    default:
+      return failure();
+    }
+
+    auto loc = op.getLoc();
+    auto genericOp =
+        rewriter.create<GenericAtomicRMWOp>(loc, op.memref(), op.indices());
+    OpBuilder bodyBuilder = OpBuilder::atBlockEnd(genericOp.getBody());
+
+    Value lhs = genericOp.getCurrentValue();
+    Value rhs = op.value();
+    Value cmp = bodyBuilder.create<CmpFOp>(loc, predicate, lhs, rhs);
+    Value select = bodyBuilder.create<SelectOp>(loc, cmp, lhs, rhs);
+    bodyBuilder.create<AtomicYieldOp>(loc, select);
+
+    rewriter.replaceOp(op, genericOp.getResult());
+    return success();
+  }
+};
+
+struct ExpandAtomic : public ExpandAtomicBase<ExpandAtomic> {
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    patterns.insert<AtomicRMWOpConverter>(&getContext());
+
+    ConversionTarget target(getContext());
+    target.addLegalOp<GenericAtomicRMWOp>();
+    target.addDynamicallyLegalOp<AtomicRMWOp>([](AtomicRMWOp op) {
+      return op.kind() != AtomicRMWKind::maxf &&
+             op.kind() != AtomicRMWKind::minf;
+    });
+    if (failed(mlir::applyPartialConversion(getFunction(), target, patterns)))
+      signalPassFailure();
+  }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> mlir::createExpandAtomicPass() {
+  return std::make_unique<ExpandAtomic>();
+}

mlir/lib/Dialect/StandardOps/Transforms/PassDetail.h

+//===- PassDetail.h - GPU Pass class details --------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef DIALECT_STANDARD_TRANSFORMS_PASSDETAIL_H_
+#define DIALECT_STANDARD_TRANSFORMS_PASSDETAIL_H_
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+class AtomicRMWOp;
+
+#define GEN_PASS_CLASSES
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h.inc"
+
+} // end namespace mlir
+
+#endif // DIALECT_STANDARD_TRANSFORMS_PASSDETAIL_H_

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

func @atomic_rmw(%I : memref<10xi32>, %ival : i32, %F : memref<10xf32>, %fval : f32, %i : index) {
   // CHECK: llvm.atomicrmw umin %{{.*}}, %{{.*}} acq_rel
   atomic_rmw "addf" %fval, %F[%i] : (f32, memref<10xf32>) -> f32
   // CHECK: llvm.atomicrmw fadd %{{.*}}, %{{.*}} acq_rel
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @cmpxchg
-func @cmpxchg(%F : memref<10xf32>, %fval : f32, %i : index) -> f32 {
-  %x = atomic_rmw "maxf" %fval, %F[%i] : (f32, memref<10xf32>) -> f32
-  // CHECK: %[[init:.*]] = llvm.load %{{.*}} : !llvm<"float*">
-  // CHECK-NEXT: llvm.br ^bb1(%[[init]] : !llvm.float)
-  // CHECK-NEXT: ^bb1(%[[loaded:.*]]: !llvm.float):
-  // CHECK-NEXT: %[[cmp:.*]] = llvm.fcmp "ogt" %[[loaded]], %{{.*}} : !llvm.float
-  // CHECK-NEXT: %[[max:.*]] = llvm.select %[[cmp]], %[[loaded]], %{{.*}} : !llvm.i1, !llvm.float
-  // CHECK-NEXT: %[[pair:.*]] = llvm.cmpxchg %{{.*}}, %[[loaded]], %[[max]] acq_rel monotonic : !llvm.float
-  // CHECK-NEXT: %[[new:.*]] = llvm.extractvalue %[[pair]][0] : !llvm<"{ float, i1 }">
-  // CHECK-NEXT: %[[ok:.*]] = llvm.extractvalue %[[pair]][1] : !llvm<"{ float, i1 }">
-  // CHECK-NEXT: llvm.cond_br %[[ok]], ^bb2, ^bb1(%[[new]] : !llvm.float)
-  // CHECK-NEXT: ^bb2:
-  return %x : f32
-  // CHECK-NEXT: llvm.return %[[new]]
-}
-
-// -----
-
 // CHECK-LABEL: func @generic_atomic_rmw
 func @generic_atomic_rmw(%I : memref<10xf32>, %i : index) -> f32 {
   %x = generic_atomic_rmw %I[%i] : memref<10xf32> {
     ^bb0(%old_value : f32):
       %c1 = constant 1.0 : f32
       atomic_yield %c1 : f32
   }
   // CHECK: [[init:%.*]] = llvm.load %{{.*}} : !llvm<"float*">

mlir/test/Dialect/Standard/expand-atomic.mlir

+// RUN: mlir-opt %s -expand-atomic -split-input-file | FileCheck %s --dump-input-on-failure
+
+// CHECK-LABEL: func @atomic_rmw_to_generic
+// CHECK-SAME: ([[F:%.*]]: memref<10xf32>, [[f:%.*]]: f32, [[i:%.*]]: index)
+func @atomic_rmw_to_generic(%F: memref<10xf32>, %f: f32, %i: index) -> f32 {
+  %x = atomic_rmw "maxf" %f, %F[%i] : (f32, memref<10xf32>) -> f32
+  return %x : f32
+}
+// CHECK: %0 = std.generic_atomic_rmw %arg0[%arg2] : memref<10xf32> {
+// CHECK: ^bb0([[CUR_VAL:%.*]]: f32):
+// CHECK:   [[CMP:%.*]] = cmpf "ogt", [[CUR_VAL]], [[f]] : f32
+// CHECK:   [[SELECT:%.*]] = select [[CMP]], [[CUR_VAL]], [[f]] : f32
+// CHECK:   atomic_yield [[SELECT]] : f32
+// CHECK: }
+// CHECK: return %0 : f32
+
+// -----
+
+// CHECK-LABEL: func @atomic_rmw_no_conversion
+func @atomic_rmw_no_conversion(%F: memref<10xf32>, %f: f32, %i: index) -> f32 {
+  %x = atomic_rmw "addf" %f, %F[%i] : (f32, memref<10xf32>) -> f32
+  return %x : f32
+}
+// CHECK-NOT: generic_atomic_rmw