[mlir][gpu] Add a pattern for transforming gpu.global_id to thread + blockId * blockDim

[fmorac]

Summary

This patch implements a rewrite pattern for transforming gpu.global_id x
to gpu.thread_id + gpu.block_id * gpu.block_dim.

[fmorac]

ping for review

[fmorac]

Rebasing main.

[makslevental]

Looks good to me (I'm surprised there isn't a general purpose utility like this that's akin to affine:detensorize).

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

 
 /// 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 GlobalIdOp op within the GPU dialect.
+void populateGpuGlobalIdPatterns(RewritePatternSet &patterns);
+
 /// 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);
+  populateGpuGlobalIdPatterns(patterns);
 }
 
 namespace gpu {
 /// Base pass class to serialize kernel functions through LLVM into
 /// user-specified IR and add the resulting blob as module attribute.
 class SerializeToBlobPass : public OperationPass<gpu::GPUModuleOp> {
 public:
   SerializeToBlobPass(TypeID passID);

mlir/lib/Dialect/GPU/CMakeLists.txt

add_mlir_dialect_library(MLIRGPUOps
   MLIRMemRefDialect
   MLIRSideEffectInterfaces
   MLIRSupport
   )
 
 add_mlir_dialect_library(MLIRGPUTransforms
   Transforms/AllReduceLowering.cpp
   Transforms/AsyncRegionRewriter.cpp
+  Transforms/GlobalIdRewriter.cpp
   Transforms/KernelOutlining.cpp
   Transforms/MemoryPromotion.cpp
   Transforms/ParallelLoopMapper.cpp
   Transforms/SerializeToBlob.cpp
   Transforms/SerializeToCubin.cpp
   Transforms/SerializeToHsaco.cpp
 
   ADDITIONAL_HEADER_DIRS
add_mlir_dialect_library(MLIRGPUTransforms
   MLIRAffineUtils
   MLIRArithDialect
   MLIRAsyncDialect
   MLIRBuiltinToLLVMIRTranslation
   MLIRDataLayoutInterfaces
   MLIRExecutionEngineUtils
   MLIRGPUOps
   MLIRIR
+  MLIRIndexDialect
   MLIRLLVMDialect
   MLIRGPUToLLVMIRTranslation
   MLIRLLVMToLLVMIRTranslation
   MLIRMemRefDialect
   MLIRPass
   MLIRSCFDialect
   MLIRSideEffectInterfaces
   MLIRSupport

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

+//===- GlobalIdRewriter.cpp - Implementation of GlobalId rewriting  -------===//
+//
+// 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 in-dialect rewriting of the global_id op for archs
+// where global_id.x = threadId.x + blockId.x * blockDim.x
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/GPU/Transforms/Passes.h"
+#include "mlir/Dialect/Index/IR/IndexOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+struct GpuGlobalIdRewriter : public OpRewritePattern<gpu::GlobalIdOp> {
+  using OpRewritePattern<gpu::GlobalIdOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(gpu::GlobalIdOp op,
+                                PatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto dim = op.getDimension();
+    auto blockId = rewriter.create<gpu::BlockIdOp>(loc, dim);
+    auto blockDim = rewriter.create<gpu::BlockDimOp>(loc, dim);
+    // Compute blockId.x * blockDim.x
+    auto tmp = rewriter.create<index::MulOp>(op.getLoc(), blockId, blockDim);
+    auto threadId = rewriter.create<gpu::ThreadIdOp>(loc, dim);
+    // Compute threadId.x + blockId.x * blockDim.x
+    rewriter.replaceOpWithNewOp<index::AddOp>(op, threadId, tmp);
+    return success();
+  }
+};
+} // namespace
+
+void mlir::populateGpuGlobalIdPatterns(RewritePatternSet &patterns) {
+  patterns.add<GpuGlobalIdRewriter>(patterns.getContext());
+}

mlir/test/Dialect/GPU/globalId-rewrite.mlir

+// RUN: mlir-opt --test-gpu-rewrite -split-input-file %s | FileCheck %s
+
+module {
+  // CHECK-LABEL: func.func @globalId
+  // CHECK-SAME: (%[[SZ:.*]]: index, %[[MEM:.*]]: memref<index, 1>) {
+  func.func @globalId(%sz : index, %mem: memref<index, 1>) {
+    gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %sz, %grid_y = %sz, %grid_z = %sz)
+               threads(%tx, %ty, %tz) in (%block_x = %sz, %block_y = %sz, %block_z = %sz) {
+      // CHECK: %[[BIDY:.*]] = gpu.block_id x
+      // CHECK-NEXT: %[[BDIMY:.*]] = gpu.block_dim x
+      // CHECK-NEXT: %[[TMPY:.*]] = index.mul %[[BIDY]], %[[BDIMY]]
+      // CHECK-NEXT: %[[TIDX:.*]] = gpu.thread_id x
+      // CHECK-NEXT: %[[GIDX:.*]] = index.add %[[TIDX]], %[[TMPY]]
+      %idx = gpu.global_id x
+      // CHECK: memref.store %[[GIDX]], %[[MEM]][] : memref<index, 1>
+      memref.store %idx, %mem[] : memref<index, 1>
+  
+      // CHECK: %[[BIDY:.*]] = gpu.block_id y
+      // CHECK-NEXT: %[[BDIMY:.*]] = gpu.block_dim y
+      // CHECK-NEXT: %[[TMPY:.*]] = index.mul %[[BIDY]], %[[BDIMY]]
+      // CHECK-NEXT: %[[TIDY:.*]] = gpu.thread_id y
+      // CHECK-NEXT: %[[GIDY:.*]] = index.add %[[TIDY]], %[[TMPY]]
+      %idy = gpu.global_id y
+      // CHECK: memref.store %[[GIDY]], %[[MEM]][] : memref<index, 1>
+      memref.store %idy, %mem[] : memref<index, 1>
+  
+      // CHECK: %[[BIDZ:.*]] = gpu.block_id z
+      // CHECK-NEXT: %[[BDIMZ:.*]] = gpu.block_dim z
+      // CHECK-NEXT: %[[TMPZ:.*]] = index.mul %[[BIDZ]], %[[BDIMZ]]
+      // CHECK-NEXT: %[[TIDZ:.*]] = gpu.thread_id z
+      // CHECK-NEXT: %[[GIDZ:.*]] = index.add %[[TIDZ]], %[[TMPZ]]
+      %idz = gpu.global_id z
+      // CHECK: memref.store %[[GIDZ]], %[[MEM]][] : memref<index, 1>
+      memref.store %idz, %mem[] : memref<index, 1>
+      gpu.terminator
+    }
+    return
+  }
+}

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

 //===- TestAllReduceLowering.cpp - Test gpu.all_reduce lowering -----------===//
 //
 // 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 lowering the gpu.all_reduce op.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
+#include "mlir/Dialect/Index/IR/IndexDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 struct TestGpuRewritePass
     : public PassWrapper<TestGpuRewritePass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestGpuRewritePass)
 
   void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<arith::ArithDialect, func::FuncDialect,
+    registry.insert<arith::ArithDialect, func::FuncDialect, index::IndexDialect,
                     memref::MemRefDialect>();
   }
   StringRef getArgument() const final { return "test-gpu-rewrite"; }
   StringRef getDescription() const final {
     return "Applies all rewrite patterns within the GPU dialect.";
   }
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());