[MLIR][GPU] Expose GpuParallelLoopMapping as non-test pass.

[csigg]

Comment fix.

[herhut]

If it useful beyond testing, I don't see why it should not be a regular pass.

[bondhugula]

LGTM - should have definitely been a regular pass. Some minor comments.

mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp

146

I think you no longer need this macro.

155

You are already in the mlir namespace.

[csigg]

Clean up namespacing.

[csigg]

mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp

146

I think you still need it for passes with internal linkage: https://mlir.llvm.org/doxygen/classmlir_1_1TypeID.html#details

155

It's a namespace-qualified definition, so the return types need the namespace qualifier as well.

I did a pass to make the namespace use more consistent with the rest of mlir.

[rriddle]

mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp

146

You don't need it for passes defined in ODS, it's already handled for you.

[bondhugula]

mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp

26–27

It was fine the way it was earlier -- we shouldn't be enclosing all of this in the mlir namespace, but instead just be defining the methods as LogicalResult mlir::gpu::.... See https://llvm.org/docs/CodingStandards.html#use-namespace-qualifiers-to-implement-previously-declared-functions
It allows one to catch build errors early.

59

Again, these should just go into the anonymous namespace.

[csigg]

Remove TYPE_ID definition.

[csigg]

mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp

26–27

These were defined in the mlir::gpu namespace before, which goes against the coding standard you linked. I moved the definition to the mlir namespace and qualified it with gpu to get early build errors. Fully qualifying them with`mlir::gpu` doesn't give us anything extra.

59

I put it in the mlir::gpu namespace because MappingLevel is a sufficiently generic term that there is risk that someone will define it in the mlir or global namespace.

146

Thanks, I didn't know.

[rriddle]

mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp

26–27

Please drop these in favor of using namespace mlir/etc. and anonymous namespaces.

mlir/include/mlir/Dialect/GPU/ParallelLoopMapper.h

 /// Sets the mapping attribute of a scf.parallel operation. Verifies that the
 /// mapping passed is valid.
 /// - the number of DimMapperAttr provided is same as the number of loops of
 ///   the `ploopOp`.
 /// - the mapping does not map multiple loops to the same processor.
 LogicalResult setMappingAttr(scf::ParallelOp ploopOp,
                              ArrayRef<ParallelLoopDimMapping> mapping);
 } // namespace gpu
-
-/// Maps the parallel loops found in the given function to workgroups. The first
-/// loop encountered will be mapped to the global workgroup and the second loop
-/// encountered to the local workgroup. Within each mapping, the first three
-/// dimensions are mapped to x/y/z hardware ids and all following dimensions are
-/// mapped to sequential loops.
-void greedilyMapParallelSCFToGPU(Region &region);
-
 } // namespace mlir
 #endif // MLIR_DIALECT_GPU_PARALLELLOOPMAPPER_H

mlir/include/mlir/Dialect/GPU/Passes.h

 /// Replaces `gpu.launch` with `gpu.launch_func` by moving the region into
 /// a separate kernel function.
 std::unique_ptr<OperationPass<ModuleOp>>
 createGpuKernelOutliningPass(StringRef dataLayoutStr = StringRef());
 
 /// Rewrites a function region so that GPU ops execute asynchronously.
 std::unique_ptr<OperationPass<func::FuncOp>> createGpuAsyncRegionPass();
 
+/// Maps the parallel loops found in the given function to workgroups. The first
+/// loop encountered will be mapped to the global workgroup and the second loop
+/// encountered to the local workgroup. Within each mapping, the first three
+/// dimensions are mapped to x/y/z hardware ids and all following dimensions are
+/// mapped to sequential loops.
+std::unique_ptr<OperationPass<func::FuncOp>> createGpuMapParallelLoopsPass();
+
 /// Collect a set of patterns to rewrite all-reduce ops within the GPU dialect.
 void populateGpuAllReducePatterns(RewritePatternSet &patterns);
 
 /// Collect all patterns to rewrite ops within the GPU dialect.
 inline void populateGpuRewritePatterns(RewritePatternSet &patterns) {
   populateGpuAllReducePatterns(patterns);
 }
 

mlir/include/mlir/Dialect/GPU/Passes.td

 }
 
 def GpuAsyncRegionPass : Pass<"gpu-async-region", "func::FuncOp"> {
   let summary = "Make GPU ops async";
   let constructor = "mlir::createGpuAsyncRegionPass()";
   let dependentDialects = ["async::AsyncDialect"];
 }
 
+def GpuMapParallelLoopsPass
+    : Pass<"gpu-map-parallel-loops", "mlir::func::FuncOp"> {
+  let summary = "Greedily maps loops to GPU hardware dimensions.";
+  let constructor = "mlir::createGpuMapParallelLoopsPass()";
+  let description = "Greedily maps loops to GPU hardware dimensions.";
+}
+
 #endif // MLIR_DIALECT_GPU_PASSES

mlir/lib/Dialect/GPU/Transforms/ParallelLoopMapper.cpp

 //===- ParallelLoopMapper.cpp - Utilities for mapping parallel loops to GPU =//
 //
 // 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 utilities to generate mappings for parallel loops to
 // GPU devices.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/GPU/ParallelLoopMapper.h"
 
+#include "PassDetail.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/Pass/Pass.h"
 
-using namespace mlir;
-using namespace mlir::gpu;
-using namespace mlir::scf;
-
 #include "mlir/Dialect/GPU/ParallelLoopMapperAttr.cpp.inc"
 #include "mlir/Dialect/GPU/ParallelLoopMapperEnums.cpp.inc"
+
 namespace mlir {
-namespace gpu {
 
-StringRef getMappingAttrName() { return "mapping"; }
+using scf::ParallelOp;
 
-ParallelLoopDimMapping getParallelLoopDimMappingAttr(Processor processor,
-                                                     AffineMap map,
-                                                     AffineMap bound) {
+StringRef gpu::getMappingAttrName() { return "mapping"; }
+
+gpu::ParallelLoopDimMapping
+gpu::getParallelLoopDimMappingAttr(Processor processor, AffineMap map,
   MLIRContext *context = map.getContext();
   OpBuilder builder(context);
   return ParallelLoopDimMapping::get(
       ProcessorAttr::get(builder.getContext(), processor),
       AffineMapAttr::get(map), AffineMapAttr::get(bound), context);
 }
 
-LogicalResult setMappingAttr(scf::ParallelOp ploopOp,
+LogicalResult gpu::setMappingAttr(ParallelOp ploopOp,
                                   ArrayRef<ParallelLoopDimMapping> mapping) {
   // Verify that each processor is mapped to only once.
   llvm::DenseSet<gpu::Processor> specifiedMappings;
   for (auto dimAttr : mapping) {
     gpu::Processor processor = getProcessor(dimAttr);
     if (processor != gpu::Processor::Sequential &&
         specifiedMappings.count(processor))
       return ploopOp.emitError(
           "invalid mapping multiple loops to same processor");
   }
   ArrayRef<Attribute> mappingAsAttrs(mapping.data(), mapping.size());
   ploopOp->setAttr(getMappingAttrName(),
                    ArrayAttr::get(ploopOp.getContext(), mappingAsAttrs));
   return success();
 }
-} // namespace gpu
-} // namespace mlir
 
+namespace gpu {
 namespace {
-
 enum MappingLevel { MapGrid = 0, MapBlock = 1, Sequential = 2 };
+} // namespace
 
 static constexpr int kNumHardwareIds = 3;
 
-} // namespace
-
 /// Bounded increment on MappingLevel. Increments to the next
 /// level unless Sequential was already reached.
-MappingLevel &operator++(MappingLevel &mappingLevel) {
+static MappingLevel &operator++(MappingLevel &mappingLevel) {
   if (mappingLevel < Sequential) {
     mappingLevel = static_cast<MappingLevel>(mappingLevel + 1);
   }
   return mappingLevel;
 }
 
 /// Computed the hardware id to use for a given mapping level. Will
 /// assign x,y and z hardware ids for the first 3 dimensions and use
 /// sequential after.
 /// TODO: Make this use x for the inner-most loop that is
 /// distributed to map to x, the next innermost to y and the next innermost to
 /// z.
-static gpu::Processor getHardwareIdForMapping(MappingLevel level,
-                                              int dimension) {
+static Processor getHardwareIdForMapping(MappingLevel level, int dimension) {
 
   if (dimension >= kNumHardwareIds || level == Sequential)
     return Processor::Sequential;
   switch (level) {
   case MapGrid:
     switch (dimension) {
     case 0:
       return Processor::BlockX;
static void mapParallelOp(ParallelOp parallelOp,
   // Parallel loop operations are immediately nested, so do not use
   // walk but just iterate over the operations.
   for (Operation &op : *parallelOp.getBody()) {
     if (ParallelOp nested = dyn_cast<ParallelOp>(op))
       mapParallelOp(nested, mappingLevel);
   }
 }
 
-void mlir::greedilyMapParallelSCFToGPU(Region &region) {
+namespace {
+struct GpuMapParallelLoopsPass
+    : public GpuMapParallelLoopsPassBase<GpuMapParallelLoopsPass> {
+  void runOnOperation() override {
+    for (Region &region : getOperation()->getRegions()) {
       region.walk([](ParallelOp parallelOp) { mapParallelOp(parallelOp); });
     }
+  }
+};
+
+} // namespace
+} // namespace gpu
+} // namespace mlir
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+mlir::createGpuMapParallelLoopsPass() {
+  return std::make_unique<gpu::GpuMapParallelLoopsPass>();
+}

mlir/test/Dialect/GPU/mapping.mlir

-// RUN: mlir-opt -test-gpu-greedy-parallel-loop-mapping -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -gpu-map-parallel-loops -split-input-file %s | FileCheck %s
 
 func.func @parallel_loop(%arg0 : index, %arg1 : index, %arg2 : index,
                     %arg3 : index) {
   %zero = arith.constant 0 : index
   %one = arith.constant 1 : index
   %four = arith.constant 4 : index
   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
                                           step (%four, %four)  {

mlir/test/lib/Dialect/GPU/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRGPUTestPasses
   TestConvertGPUKernelToCubin.cpp
   TestConvertGPUKernelToHsaco.cpp
   TestGpuMemoryPromotion.cpp
-  TestGpuParallelLoopMapping.cpp
   TestGpuRewrite.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRFunc
   MLIRGPUOps

mlir/test/lib/Dialect/GPU/TestGpuParallelLoopMapping.cpp

-//===- TestGPUParallelLoopMapping.cpp - Test pass for GPU loop mapping ----===//
-//
-// 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 the pass testing the utilities for mapping parallel
-// loops to gpu hardware ids.
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Dialect/GPU/ParallelLoopMapper.h"
-#include "mlir/Pass/Pass.h"
-
-using namespace mlir;
-
-namespace {
-/// Simple pass for testing the mapping of parallel loops to hardware ids using
-/// a greedy mapping strategy.
-struct TestGpuGreedyParallelLoopMappingPass
-    : public PassWrapper<TestGpuGreedyParallelLoopMappingPass,
-                         OperationPass<>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
-      TestGpuGreedyParallelLoopMappingPass)
-
-  StringRef getArgument() const final {
-    return "test-gpu-greedy-parallel-loop-mapping";
-  }
-  StringRef getDescription() const final {
-    return "Greedily maps all parallel loops to gpu hardware ids.";
-  }
-  void runOnOperation() override {
-    for (Region &region : getOperation()->getRegions())
-      greedilyMapParallelSCFToGPU(region);
-  }
-};
-} // namespace
-
-namespace mlir {
-namespace test {
-void registerTestGpuParallelLoopMappingPass() {
-  PassRegistration<TestGpuGreedyParallelLoopMappingPass>();
-}
-} // namespace test
-} // namespace mlir

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

 void registerTestDataLayoutQuery();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDiagnosticsPass();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandTanhPass();
 void registerTestComposeSubView();
 void registerTestMultiBuffering();
-void registerTestGpuParallelLoopMappingPass();
 void registerTestIRVisitorsPass();
 void registerTestGenericIRVisitorsPass();
 void registerTestGenericIRVisitorsInterruptPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgElementwiseFusion();
 void registerTestLinalgFusionTransforms();
 void registerTestLinalgTensorFusionTransforms();
#endif
   mlir::test::registerTestDecomposeCallGraphTypes();
   mlir::test::registerTestDataFlowPass();
   mlir::test::registerTestDataLayoutQuery();
   mlir::test::registerTestDominancePass();
   mlir::test::registerTestDynamicPipelinePass();
   mlir::test::registerTestExpandTanhPass();
   mlir::test::registerTestComposeSubView();
   mlir::test::registerTestMultiBuffering();
-  mlir::test::registerTestGpuParallelLoopMappingPass();
   mlir::test::registerTestIRVisitorsPass();
   mlir::test::registerTestGenericIRVisitorsPass();
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLinalgCodegenStrategy();
   mlir::test::registerTestLinalgElementwiseFusion();
   mlir::test::registerTestLinalgFusionTransforms();
   mlir::test::registerTestLinalgTensorFusionTransforms();
   mlir::test::registerTestLinalgTiledLoopFusionTransforms();