f5d83fb404
There's no world where the subgroup ID (or the intermediate values needed to compute it) will be negative or will have signed overflow. This commit adds flags accordingly, which is helpful as this is a rather low-level rewrite that might run after the analyses that would ordinarily add these flags.
120 lines
5.1 KiB
C++
120 lines
5.1 KiB
C++
//===- SubgroupIdRewriter.cpp - Implementation of SubgroupId 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 gpu.subgroup_id op for archs
|
|
// where:
|
|
// subgroup_id = (tid.x + dim.x * (tid.y + dim.y * tid.z)) / subgroup_size
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#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"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
struct GpuSubgroupIdRewriter final : OpRewritePattern<gpu::SubgroupIdOp> {
|
|
using OpRewritePattern<gpu::SubgroupIdOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(gpu::SubgroupIdOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
// Calculation of the thread's subgroup identifier.
|
|
//
|
|
// The process involves mapping the thread's 3D identifier within its
|
|
// block (b_id.x, b_id.y, b_id.z) to a 1D linear index.
|
|
// This linearization assumes a layout where the x-dimension (w_dim.x)
|
|
// varies most rapidly (i.e., it is the innermost dimension).
|
|
//
|
|
// The formula for the linearized thread index is:
|
|
// L = tid.x + dim.x * (tid.y + (dim.y * tid.z))
|
|
//
|
|
// Subsequently, the range of linearized indices [0, N_threads-1] is
|
|
// divided into consecutive, non-overlapping segments, each representing
|
|
// a subgroup of size 'subgroup_size'.
|
|
//
|
|
// Example Partitioning (N = subgroup_size):
|
|
// | Subgroup 0 | Subgroup 1 | Subgroup 2 | ... |
|
|
// | Indices 0..N-1 | Indices N..2N-1 | Indices 2N..3N-1| ... |
|
|
//
|
|
// The subgroup identifier is obtained via integer division of the
|
|
// linearized thread index by the predefined 'subgroup_size'.
|
|
//
|
|
// subgroup_id = floor( L / subgroup_size )
|
|
// = (tid.x + dim.x * (tid.y + dim.y * tid.z)) /
|
|
// subgroup_size
|
|
|
|
Location loc = op->getLoc();
|
|
Type indexType = rewriter.getIndexType();
|
|
|
|
auto asMaybeIndexAttr = [&](std::optional<uint32_t> bound) -> IntegerAttr {
|
|
if (!bound)
|
|
return IntegerAttr();
|
|
return IntegerAttr::get(
|
|
indexType, static_cast<int64_t>(static_cast<uint64_t>(*bound)));
|
|
};
|
|
|
|
IntegerAttr maybeKnownDimX =
|
|
asMaybeIndexAttr(gpu::getKnownDimensionSizeAround(
|
|
op, gpu::DimensionKind::Block, gpu::Dimension::x));
|
|
IntegerAttr maybeKnownDimY =
|
|
asMaybeIndexAttr(gpu::getKnownDimensionSizeAround(
|
|
op, gpu::DimensionKind::Block, gpu::Dimension::y));
|
|
IntegerAttr maybeKnownDimZ =
|
|
asMaybeIndexAttr(gpu::getKnownDimensionSizeAround(
|
|
op, gpu::DimensionKind::Block, gpu::Dimension::z));
|
|
|
|
Value dimX, dimY;
|
|
if (maybeKnownDimX)
|
|
dimX = arith::ConstantOp::create(rewriter, loc, maybeKnownDimX);
|
|
else
|
|
dimX = gpu::BlockDimOp::create(rewriter, loc, gpu::Dimension::x);
|
|
if (maybeKnownDimY)
|
|
dimY = arith::ConstantOp::create(rewriter, loc, maybeKnownDimY);
|
|
else
|
|
dimY = gpu::BlockDimOp::create(rewriter, loc, gpu::Dimension::y);
|
|
|
|
Value tidX = gpu::ThreadIdOp::create(rewriter, loc, gpu::Dimension::x,
|
|
maybeKnownDimX);
|
|
Value tidY = gpu::ThreadIdOp::create(rewriter, loc, gpu::Dimension::y,
|
|
maybeKnownDimY);
|
|
Value tidZ = gpu::ThreadIdOp::create(rewriter, loc, gpu::Dimension::z,
|
|
maybeKnownDimZ);
|
|
|
|
// Block dimensions don't exceed a signed int32_t maximum, and neither does
|
|
// their product, on any realistic hardware, nor would any targets compile
|
|
// with index < 32 bits, so we can assert no overflow.
|
|
auto flags =
|
|
arith::IntegerOverflowFlags::nsw | arith::IntegerOverflowFlags::nuw;
|
|
Value dimYxIdZ =
|
|
arith::MulIOp::create(rewriter, loc, indexType, dimY, tidZ, flags);
|
|
Value dimYxIdZPlusIdY =
|
|
arith::AddIOp::create(rewriter, loc, indexType, dimYxIdZ, tidY, flags);
|
|
Value dimYxIdZPlusIdYTimesDimX = arith::MulIOp::create(
|
|
rewriter, loc, indexType, dimX, dimYxIdZPlusIdY, flags);
|
|
Value idXPlusDimYxIdZPlusIdYTimesDimX = arith::AddIOp::create(
|
|
rewriter, loc, indexType, tidX, dimYxIdZPlusIdYTimesDimX, flags);
|
|
Value subgroupSize = gpu::SubgroupSizeOp::create(
|
|
rewriter, loc, rewriter.getIndexType(), /*upper_bound = */ nullptr);
|
|
Value subgroupIdOp =
|
|
arith::DivUIOp::create(rewriter, loc, indexType,
|
|
idXPlusDimYxIdZPlusIdYTimesDimX, subgroupSize);
|
|
rewriter.replaceOp(op, {subgroupIdOp});
|
|
return success();
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
void mlir::populateGpuSubgroupIdPatterns(RewritePatternSet &patterns) {
|
|
patterns.add<GpuSubgroupIdRewriter>(patterns.getContext());
|
|
}
|