[mlir] Move `expand-atomic` to test passes. Enable the pattern in Std2LLVM rewrite.

[pifon2a]

Summary

This is consistent with the patterns in ExpandTanh.cpp and
ExpandMemRefReshape.cpp. The application of this pattern should happen in
StandardToLLVM lowering not in a separate pass.

[pifon2a]

Fix comments.

[ftynse]

I'm not convinced about this one. I have not reviewed https://reviews.llvm.org/D90235 that added the memref reshape expansion, and I would have complained about that too. TanhExpansion doesn't seem to be plugged by default in the lowering.

My concerns are:

• This increases the complexity of the std-to-llvm conversion, making it also apply some std-to-std patterns. (At some point, we actually had a direct conversion from std to LLVM IR, but decided to add the LLVM dialect to make conversion more progressive.) It's unclear where to draw the line on adding such patterns: we will also include the ceildiv/floordiv expansion proposed in another review thread, we could also need the expansion of pointwise tensor operations into loops, then cfg (be they std.addf on tensor operands, or the new operation being discussed on the forum), and so on, further increasing complexity. At the same time, std-to-llvm conversion is already big and complex enough, with StandardToLLVM.cpp being the second largest file in LoC after StandardDialect/Ops.cpp in the codebase...
• I can imagine users that only want the conversion and prefer to handle expansions differently, by targeting LLVM intrinsics or specific instructions if they go to lower levels without leaving MLIR. Different platforms may support different atomics. This was the argument for keeping tanh expansion outside the conversion to LLVM as far as I recall.
• The layering looks weird in https://reviews.llvm.org/D90235, with added dependency on from StdToLLVMConversion on StdTransforms. These things look like they should be independent.

I think I would be in favor of some expansion pass (just don't call it -std-to-std!) that eliminates operations we don't want to see at the entry of the conversion pass, similarly to the pass we have that prepares the LLVM dialect for translation to LLVM IR because the dialect is slightly more relaxed. While this increases the complexity of the pipeline by having one extra pass, the decrease in complexity and semantic load of each pass largely compensates that IMO.

mlir/test/lib/Transforms/TestExpandAtomic.cpp

1

Not: the description is copy-pasta

[pifon2a]

I can imagine users that only want the conversion and prefer to handle expansions differently, by targeting LLVM intrinsics or specific instructions if they go to lower levels without leaving MLIR. Different platforms may support different atomics. This was the argument for keeping tanh expansion outside the conversion to LLVM as far as I recall.

ExpandMemRefReshape and ExpandAtomic differ from ExpandTanh, because they exist only for code reuse, without them MemRefReshapeOp and AtomicRMWOp are impossible to lower to LLVM for all "flavors".

I propose the following: I will remove`populateExpand[Atomic|MemRefReshape]patterns` from StandardToLLVM.cpp, but still leave TestExpandAtomic, TestExpandMemRefReshape in mlir/test/Transforms. Then the users can add these patterns to their LLVM lowering pass if needed, or a combination of these passes to their custom "std2std" pass

[ftynse]

In D90785#2378621, @pifon2a wrote:

I propose the following: I will remove`populateExpand[Atomic|MemRefReshape]patterns` from StandardToLLVM.cpp, but still leave TestExpandAtomic, TestExpandMemRefReshape in mlir/test/Transforms. Then the users can add these patterns to their LLVM lowering pass if needed, or a combination of these passes to their custom "std2std" pass

What would be the default in-tree way of doing this? We wouldn't expect people to run -test-expand-atomic -convert-std-to-llvm or having to write a custom pass every time they need it, hence my suggestion of an additional "legalization" pass.

[pifon2a]

abandoned in favor of https://reviews.llvm.org/D91082

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

 
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Bufferize.h"
 
 namespace mlir {
 
 class OwningRewritePatternList;
 
-/// Creates an instance of the ExpandAtomic pass.
-std::unique_ptr<Pass> createExpandAtomicPass();
+void populateExpandAtomicPattern(OwningRewritePatternList &patterns,
+                                 MLIRContext *ctx);
 
 void populateExpandMemRefReshapePattern(OwningRewritePatternList &patterns,
                                         MLIRContext *ctx);
 
 void populateExpandTanhPattern(OwningRewritePatternList &patterns,
                                MLIRContext *ctx);
 
 void populateStdBufferizePatterns(MLIRContext *context,

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()";
-}
-
 def StdBufferize : FunctionPass<"std-bufferize"> {
   let summary = "Bufferize the std dialect";
   let constructor = "mlir::createStdBufferizePass()";
   let dependentDialects = ["scf::SCFDialect"];
 }
 
 def FuncBufferize : Pass<"func-bufferize", "ModuleOp"> {
   let summary = "Bufferize func/call/return ops";

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

patterns.insert<
       TruncateIOpLowering,
       UIToFPLowering,
       UnsignedDivIOpLowering,
       UnsignedRemIOpLowering,
       UnsignedShiftRightOpLowering,
       XOrOpLowering,
       ZeroExtendIOpLowering>(converter);
   // clang-format on
+  populateExpandAtomicPattern(patterns, &converter.getContext());
 }
 
 void mlir::populateStdToLLVMMemoryConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   // clang-format off
   patterns.insert<
       AssumeAlignmentOpLowering,
       DeallocOpLowering,

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

LogicalResult matchAndRewrite(AtomicRMWOp op,
     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,
-                                            std::move(patterns))))
-      signalPassFailure();
-  }
-};
-
 } // namespace
 
-std::unique_ptr<Pass> mlir::createExpandAtomicPass() {
-  return std::make_unique<ExpandAtomic>();
+void mlir::populateExpandAtomicPattern(OwningRewritePatternList &patterns,
+                                       MLIRContext *ctx) {
+  patterns.insert<AtomicRMWOpConverter>(ctx);
 }

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

-// RUN: mlir-opt %s -expand-atomic -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-expand-atomic -split-input-file | FileCheck %s
 
 // 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> {

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestAffineLoopParametricTiling.cpp
+  TestExpandAtomic.cpp
   TestExpandMemRefReshape.cpp
   TestExpandTanh.cpp
   TestCallGraph.cpp
   TestConstantFold.cpp
   TestConvVectorization.cpp
   TestConvertCallOp.cpp
   TestConvertGPUKernelToCubin.cpp
   TestConvertGPUKernelToHsaco.cpp

mlir/test/lib/Transforms/TestExpandAtomic.cpp

+//===- TestExpandAtomic.cpp - Test expand tanh op into exp form ------===//
+//
+// 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 contains test passes for expanding atomic.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+struct TestExpandAtomicPass
+    : public PassWrapper<TestExpandAtomicPass, FunctionPass> {
+  void runOnFunction() override;
+};
+} // end anonymous namespace
+
+void TestExpandAtomicPass::runOnFunction() {
+  OwningRewritePatternList patterns;
+  populateExpandAtomicPattern(patterns, &getContext());
+  applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+}
+
+namespace mlir {
+void registerTestExpandAtomicPass() {
+  PassRegistration<TestExpandAtomicPass> pass("test-expand-atomic",
+                                              "Test expanding atomic");
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvVectorization();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestConvertGPUKernelToHsacoPass();
 void registerTestDominancePass();
 void registerTestDialect(DialectRegistry &);
 void registerTestDynamicPipelinePass();
+void registerTestExpandAtomicPass();
 void registerTestExpandMemRefReshapePass();
 void registerTestExpandTanhPass();
 void registerTestFinalizingBufferizePass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();

 #if MLIR_ROCM_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToHsacoPass();
 #endif
   registerTestAffineLoopParametricTilingPass();
   registerTestDominancePass();
   registerTestDynamicPipelinePass();
   registerTestFinalizingBufferizePass();
   registerTestFunc();
-  registerTestExpandTanhPass();
+  registerTestExpandAtomicPass();
   registerTestExpandMemRefReshapePass();
+  registerTestExpandTanhPass();
   registerTestGpuMemoryPromotionPass();
   registerTestInterfaces();
   registerTestLinalgCodegenStrategy();
   registerTestLinalgFusionTransforms();
   registerTestLinalgGreedyFusion();
   registerTestLinalgHoisting();
   registerTestLinalgTransforms();
   registerTestLivenessPass();