1//===- SubgroupIdRewriter.cpp - Implementation of SubgroupId rewriting ----===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file implements in-dialect rewriting of the gpu.subgroup_id op for archs
10// where:
11// subgroup_id = (tid.x + dim.x * (tid.y + dim.y * tid.z)) / subgroup_size
12//
13//===----------------------------------------------------------------------===//
14
15#include "mlir/Dialect/GPU/IR/GPUDialect.h"
16#include "mlir/Dialect/GPU/Transforms/Passes.h"
17#include "mlir/Dialect/Index/IR/IndexOps.h"
18#include "mlir/IR/Builders.h"
19#include "mlir/IR/PatternMatch.h"
20
21using namespace mlir;
22
23namespace {
24struct GpuSubgroupIdRewriter final : OpRewritePattern<gpu::SubgroupIdOp> {
25 using OpRewritePattern<gpu::SubgroupIdOp>::OpRewritePattern;
26
27 LogicalResult matchAndRewrite(gpu::SubgroupIdOp op,
28 PatternRewriter &rewriter) const override {
29 // Calculation of the thread's subgroup identifier.
30 //
31 // The process involves mapping the thread's 3D identifier within its
32 // block (b_id.x, b_id.y, b_id.z) to a 1D linear index.
33 // This linearization assumes a layout where the x-dimension (w_dim.x)
34 // varies most rapidly (i.e., it is the innermost dimension).
35 //
36 // The formula for the linearized thread index is:
37 // L = tid.x + dim.x * (tid.y + (dim.y * tid.z))
38 //
39 // Subsequently, the range of linearized indices [0, N_threads-1] is
40 // divided into consecutive, non-overlapping segments, each representing
41 // a subgroup of size 'subgroup_size'.
42 //
43 // Example Partitioning (N = subgroup_size):
44 // | Subgroup 0 | Subgroup 1 | Subgroup 2 | ... |
45 // | Indices 0..N-1 | Indices N..2N-1 | Indices 2N..3N-1| ... |
46 //
47 // The subgroup identifier is obtained via integer division of the
48 // linearized thread index by the predefined 'subgroup_size'.
49 //
50 // subgroup_id = floor( L / subgroup_size )
51 // = (tid.x + dim.x * (tid.y + dim.y * tid.z)) /
52 // subgroup_size
53
54 Location loc = op->getLoc();
55 Type indexType = rewriter.getIndexType();
56
57 Value dimX = rewriter.create<gpu::BlockDimOp>(location: loc, args: gpu::Dimension::x);
58 Value dimY = rewriter.create<gpu::BlockDimOp>(location: loc, args: gpu::Dimension::y);
59 Value tidX = rewriter.create<gpu::ThreadIdOp>(location: loc, args: gpu::Dimension::x);
60 Value tidY = rewriter.create<gpu::ThreadIdOp>(location: loc, args: gpu::Dimension::y);
61 Value tidZ = rewriter.create<gpu::ThreadIdOp>(location: loc, args: gpu::Dimension::z);
62
63 Value dimYxIdZ = rewriter.create<arith::MulIOp>(location: loc, args&: indexType, args&: dimY, args&: tidZ);
64 Value dimYxIdZPlusIdY =
65 rewriter.create<arith::AddIOp>(location: loc, args&: indexType, args&: dimYxIdZ, args&: tidY);
66 Value dimYxIdZPlusIdYTimesDimX =
67 rewriter.create<arith::MulIOp>(location: loc, args&: indexType, args&: dimX, args&: dimYxIdZPlusIdY);
68 Value IdXPlusDimYxIdZPlusIdYTimesDimX = rewriter.create<arith::AddIOp>(
69 location: loc, args&: indexType, args&: tidX, args&: dimYxIdZPlusIdYTimesDimX);
70 Value subgroupSize = rewriter.create<gpu::SubgroupSizeOp>(
71 location: loc, args: rewriter.getIndexType(), /*upper_bound = */ args: nullptr);
72 Value subgroupIdOp = rewriter.create<arith::DivUIOp>(
73 location: loc, args&: indexType, args&: IdXPlusDimYxIdZPlusIdYTimesDimX, args&: subgroupSize);
74 rewriter.replaceOp(op, newValues: {subgroupIdOp});
75 return success();
76 }
77};
78
79} // namespace
80
81void mlir::populateGpuSubgroupIdPatterns(RewritePatternSet &patterns) {
82 patterns.add<GpuSubgroupIdRewriter>(arg: patterns.getContext());
83}
84

source code of mlir/lib/Dialect/GPU/Transforms/SubgroupIdRewriter.cpp