42 #include "../GPUCommon/GPUOpsLowering.h"
43 #include "../GPUCommon/IndexIntrinsicsOpLowering.h"
46 #define GEN_PASS_DEF_CONVERTGPUOPSTOROCDLOPS
47 #include "mlir/Conversion/Passes.h.inc"
57 int64_t intWidth = cast<IntegerType>(value.
getType()).getWidth();
59 auto indexBitwidthType =
62 if (indexBitwidth > intWidth) {
63 return LLVM::SExtOp::create(rewriter, loc, indexBitwidthType, value);
65 if (indexBitwidth < intWidth) {
66 return LLVM::TruncOp::create(rewriter, loc, indexBitwidthType, value);
74 bool canBeBare =
true;
75 for (
Type type : func.getArgumentTypes())
76 if (
auto memrefTy = dyn_cast<BaseMemRefType>(type))
86 LLVM::LLVMDialect::getNoUndefAttrName(), rewriter.
getUnitAttr());
88 LLVM::LLVMDialect::getRangeAttrName(),
92 LLVM::LLVMDialect::getRangeAttrName(),
95 Value mbcntLo = ROCDL::MbcntLoOp::create(
96 rewriter, loc, int32Type, minus1, zero, {},
99 Value laneId = ROCDL::MbcntHiOp::create(
100 rewriter, loc, int32Type, minus1, mbcntLo, {},
106 "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32"
107 "-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:"
109 "32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:"
110 "64-S32-A5-G1-ni:7:8:9";
117 matchAndRewrite(gpu::LaneIdOp op, gpu::LaneIdOp::Adaptor adaptor,
131 const unsigned indexBitwidth = getTypeConverter()->getIndexTypeBitwidth();
132 if (indexBitwidth > 32) {
133 laneId = LLVM::SExtOp::create(
135 }
else if (indexBitwidth < 32) {
136 laneId = LLVM::TruncOp::create(
148 amdgpu::Chipset chipset)
153 matchAndRewrite(gpu::SubgroupSizeOp op, gpu::SubgroupSizeOp::Adaptor adaptor,
155 LLVM::ConstantRangeAttr bounds =
nullptr;
156 bool isBeforeGfx10 = chipset.majorVersion < 10;
157 if (
auto upperBoundAttr = op.getUpperBoundAttr()) {
158 bounds = rewriter.
getAttr<LLVM::ConstantRangeAttr>(
159 32, isBeforeGfx10 ? 64 : 32,
160 op.getUpperBoundAttr().getInt() + 1);
162 Value wavefrontOp = ROCDL::WavefrontSizeOp::create(
163 rewriter, op.getLoc(), rewriter.
getI32Type(), bounds);
165 *getTypeConverter());
170 const amdgpu::Chipset chipset;
173 static bool isSupportedReadLaneType(
Type type) {
177 isa<Float16Type, BFloat16Type, Float32Type, Float64Type,
178 LLVM::LLVMPointerType>(type);
181 struct GPUSubgroupBroadcastOpToROCDL
186 matchAndRewrite(gpu::SubgroupBroadcastOp op, OpAdaptor adaptor,
188 Value src = adaptor.getSrc();
189 if (!isSupportedReadLaneType(src.
getType()))
192 if (adaptor.getBroadcastType() == gpu::BroadcastType::specific_lane) {
223 matchAndRewrite(gpu::ShuffleOp op, OpAdaptor adaptor,
226 Value initShflValue = adaptor.getValue();
231 Value width = adaptor.getWidth();
232 Value zero = LLVM::ConstantOp::create(rewriter, loc, int32Type, 0);
233 Value negwidth = LLVM::SubOp::create(rewriter, loc, int32Type, zero, width);
234 Value add = LLVM::AddOp::create(rewriter, loc, int32Type, srcLaneId, width);
235 Value widthOrZeroIfOutside =
236 LLVM::AndOp::create(rewriter, loc, int32Type,
add, negwidth);
239 switch (op.getMode()) {
240 case gpu::ShuffleMode::UP:
241 dstLane = LLVM::SubOp::create(rewriter, loc, int32Type, srcLaneId,
242 adaptor.getOffset());
244 case gpu::ShuffleMode::DOWN:
245 dstLane = LLVM::AddOp::create(rewriter, loc, int32Type, srcLaneId,
246 adaptor.getOffset());
248 case gpu::ShuffleMode::XOR:
249 dstLane = LLVM::XOrOp::create(rewriter, loc, int32Type, srcLaneId,
250 adaptor.getOffset());
252 case gpu::ShuffleMode::IDX:
253 dstLane = adaptor.getOffset();
256 Value isActiveSrcLane = LLVM::ICmpOp::create(
257 rewriter, loc, LLVM::ICmpPredicate::slt, dstLane, widthOrZeroIfOutside);
258 Value selectDstLane = LLVM::SelectOp::create(rewriter, loc, isActiveSrcLane,
260 Value two = LLVM::ConstantOp::create(rewriter, loc, int32Type, 2);
261 Value dwordAlignedDstLane =
262 LLVM::ShlOp::create(rewriter, loc, int32Type, selectDstLane, two);
267 for (
Value v : decomposed) {
268 Value res = ROCDL::DsBpermuteOp::create(rewriter, loc, int32Type,
269 dwordAlignedDstLane, v);
270 swizzled.emplace_back(res);
274 rewriter.
replaceOp(op, {shflValue, isActiveSrcLane});
280 #include "GPUToROCDL.cpp.inc"
287 struct LowerGpuOpsToROCDLOpsPass final
288 :
public impl::ConvertGpuOpsToROCDLOpsBase<LowerGpuOpsToROCDLOpsPass> {
292 Base::getDependentDialects(registry);
296 void runOnOperation()
override {
297 gpu::GPUModuleOp m = getOperation();
300 auto llvmDataLayout = m->getAttrOfType<StringAttr>(
301 LLVM::LLVMDialect::getDataLayoutAttrName());
302 if (!llvmDataLayout) {
304 m->setAttr(LLVM::LLVMDialect::getDataLayoutAttrName(), llvmDataLayout);
307 for (
auto func : m.getOps<func::FuncOp>()) {
308 func->setAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(),
313 if (
failed(maybeChipset)) {
315 return signalPassFailure();
320 ctx,
DataLayout(cast<DataLayoutOpInterface>(m.getOperation())));
321 options.dataLayout = llvm::DataLayout(llvmDataLayout.getValue());
323 options.overrideIndexBitwidth(indexBitwidth);
325 if (useBarePtrCallConv) {
326 options.useBarePtrCallConv =
true;
328 m.walk([](gpu::GPUFuncOp func) ->
WalkResult {
335 "bare pointer calling convention requires all memrefs to "
336 "have static shape and use the identity map");
337 return signalPassFailure();
353 converter, [](gpu::AddressSpace space) {
355 case gpu::AddressSpace::Global:
357 case gpu::AddressSpace::Workgroup:
359 case gpu::AddressSpace::Private:
362 llvm_unreachable(
"unknown address space enum value");
369 llvm::SmallDenseSet<StringRef> allowedDialectsSet(allowedDialects.begin(),
370 allowedDialects.end());
372 bool allowed = allowedDialectsSet.contains(dialect->getNamespace());
374 if (!allowedDialectsSet.empty() && !allowed)
377 auto *iface = dyn_cast<ConvertToLLVMPatternInterface>(dialect);
383 <<
"dialect does not implement ConvertToLLVMPatternInterface: "
384 << dialect->getNamespace();
385 return signalPassFailure();
390 iface->populateConvertToLLVMConversionPatterns(target, converter,
402 auto reqdWorkGroupSizeAttrHelper =
403 rocdlDialect->getReqdWorkGroupSizeAttrHelper();
404 auto flatWorkGroupSizeAttrHelper =
405 rocdlDialect->getFlatWorkGroupSizeAttrHelper();
408 m.walk([&](LLVM::LLVMFuncOp op) {
409 if (reqdWorkGroupSizeAttrHelper.isAttrPresent(op)) {
410 auto blockSizes = reqdWorkGroupSizeAttrHelper.getAttr(op);
413 uint32_t flatSize = 1;
414 for (uint32_t size : blockSizes.asArrayRef()) {
417 StringAttr flatSizeAttr =
419 flatWorkGroupSizeAttrHelper.setAttr(op, flatSizeAttr);
432 target.
addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::Exp2Op, LLVM::FCeilOp,
433 LLVM::FFloorOp, LLVM::FRemOp, LLVM::LogOp, LLVM::Log10Op,
434 LLVM::Log2Op, LLVM::PowOp, LLVM::SinOp>();
437 return any_of(op->getOperandTypes(), llvm::IsaPred<Float32Type>);
440 target.
addLegalOp<gpu::YieldOp, gpu::GPUModuleOp>();
454 ROCDL::ThreadIdYOp, ROCDL::ThreadIdZOp>>(
455 converter, IndexKind::Block, IntrType::Id);
457 gpu::BlockIdOp, ROCDL::BlockIdXOp, ROCDL::BlockIdYOp, ROCDL::BlockIdZOp>>(
458 converter, IndexKind::Grid, IntrType::Id);
461 ROCDL::BlockDimYOp, ROCDL::BlockDimZOp>>(
462 converter, IndexKind::Block, IntrType::Dim);
464 gpu::GridDimOp, ROCDL::GridDimXOp, ROCDL::GridDimYOp, ROCDL::GridDimZOp>>(
465 converter, IndexKind::Grid, IntrType::Dim);
470 ROCDL::ROCDLDialect::kPrivateMemoryAddressSpace,
471 ROCDL::ROCDLDialect::kSharedMemoryAddressSpace,
472 rocdlDialect->getKernelAttrHelper().getName(),
473 rocdlDialect->getReqdWorkGroupSizeAttrHelper().getName()});
483 patterns.add<GPUShuffleOpLowering, GPULaneIdOpToROCDL,
484 GPUSubgroupBroadcastOpToROCDL>(converter);
485 patterns.add<GPUSubgroupSizeOpToROCDL>(converter, chipset);
static Value getZero(OpBuilder &b, Location loc, Type elementType)
Get zero value for an element type.
static MLIRContext * getContext(OpFoldResult val)
static Value getLaneId(RewriterBase &rewriter, Location loc)
static bool canBeCalledWithBarePointers(gpu::GPUFuncOp func)
Returns true if the given gpu.func can be safely called using the bare pointer calling convention.
static constexpr StringLiteral amdgcnDataLayout
static Value truncOrExtToLLVMType(ConversionPatternRewriter &rewriter, Location loc, Value value, const LLVMTypeConverter &converter)
static llvm::ManagedStatic< PassManagerOptions > options
MLIRContext * getContext() const
ArrayAttr getArrayAttr(ArrayRef< Attribute > value)
DictionaryAttr getDictionaryAttr(ArrayRef< NamedAttribute > value)
NamedAttribute getNamedAttr(StringRef name, Attribute val)
Attr getAttr(Args &&...args)
Get or construct an instance of the attribute Attr with provided arguments.
This class implements a pattern rewriter for use with ConversionPatterns.
void replaceOp(Operation *op, ValueRange newValues) override
Replace the given operation with the new values.
This class describes a specific conversion target.
void addLegalOp(OperationName op)
Register the given operations as legal.
void addLegalDialect(StringRef name, Names... names)
Register the operations of the given dialects as legal.
void addDynamicallyLegalOp(OperationName op, const DynamicLegalityCallbackFn &callback)
Register the given operation as dynamically legal and set the dynamic legalization callback to the on...
void addIllegalDialect(StringRef name, Names... names)
Register the operations of the given dialects as illegal, i.e.
void addIllegalOp(OperationName op)
Register the given operation as illegal, i.e.
Utility class for operation conversions targeting the LLVM dialect that match exactly one source oper...
ConvertOpToLLVMPattern(const LLVMTypeConverter &typeConverter, PatternBenefit benefit=1)
The main mechanism for performing data layout queries.
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
Dialects are groups of MLIR operations, types and attributes, as well as behavior associated with the...
Derived class that automatically populates legalization information for different LLVM ops.
Conversion from types to the LLVM IR dialect.
static bool canConvertToBarePtr(BaseMemRefType type)
Check if a memref type can be converted to a bare pointer.
MLIRContext & getContext() const
Returns the MLIR context.
unsigned getIndexTypeBitwidth() const
Gets the bitwidth of the index type when converted to LLVM.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Options to control the LLVM lowering.
MLIRContext is the top-level object for a collection of MLIR operations.
Dialect * getLoadedDialect(StringRef name)
Get a registered IR dialect with the given namespace.
std::vector< Dialect * > getLoadedDialects()
Return information about all IR dialects loaded in the context.
NamedAttribute represents a combination of a name and an Attribute value.
Operation is the basic unit of execution within MLIR.
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
bool isInteger() const
Return true if this is an integer type (with the specified width).
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
A utility result that is used to signal how to proceed with an ongoing walk:
static WalkResult advance()
bool wasInterrupted() const
Returns true if the walk was interrupted.
static WalkResult interrupt()
static ConstantIntOp create(OpBuilder &builder, Location location, int64_t value, unsigned width)
Value composeValue(OpBuilder &builder, Location loc, ValueRange src, Type dstType)
Composes a set of src values into a single value of type dstType through series of bitcasts and vecto...
SmallVector< Value > decomposeValue(OpBuilder &builder, Location loc, Value src, Type dstType)
Decomposes a src value into a set of values of type dstType through series of bitcasts and vector ops...
Runtime
Potential runtimes for AMD GPU kernels.
Include the generated interface declarations.
void populateGpuToROCDLConversionPatterns(const LLVMTypeConverter &converter, RewritePatternSet &patterns, gpu::amd::Runtime runtime, amdgpu::Chipset chipset)
Collect a set of patterns to convert from the GPU dialect to ROCDL.
static constexpr unsigned kDeriveIndexBitwidthFromDataLayout
Value to pass as bitwidth for the index type when the converter is expected to derive the bitwidth fr...
LogicalResult applyPatternsGreedily(Region ®ion, const FrozenRewritePatternSet &patterns, GreedyRewriteConfig config=GreedyRewriteConfig(), bool *changed=nullptr)
Rewrite ops in the given region, which must be isolated from above, by repeatedly applying the highes...
void populateGpuRewritePatterns(RewritePatternSet &patterns)
Collect all patterns to rewrite ops within the GPU dialect.
InFlightDiagnostic emitError(Location loc)
Utility method to emit an error message using this location.
void configureGpuToROCDLConversionLegality(ConversionTarget &target)
Configure target to convert from the GPU dialect to ROCDL.
const FrozenRewritePatternSet & patterns
void registerConvertToLLVMDependentDialectLoading(DialectRegistry ®istry)
Register the extension that will load dependent dialects for LLVM conversion.
void populateGpuMemorySpaceAttributeConversions(TypeConverter &typeConverter, const MemorySpaceMapping &mapping)
Populates memory space attribute conversion rules for lowering gpu.address_space to integer values.
void populateAMDGPUToROCDLConversionPatterns(LLVMTypeConverter &converter, RewritePatternSet &patterns, amdgpu::Chipset chipset)
Note: This function will also add conversions for the AMDGPU-specific address spaces,...
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
void populateGpuPromoteShuffleToAMDGPUPatterns(RewritePatternSet &patterns, std::optional< amdgpu::Chipset > maybeChipset)
Tries to promote gpu.shuffles to specialized AMDGPU intrinsics.
LogicalResult applyPartialConversion(ArrayRef< Operation * > ops, const ConversionTarget &target, const FrozenRewritePatternSet &patterns, ConversionConfig config=ConversionConfig())
Below we define several entry points for operation conversion.
void populateMathToROCDLConversionPatterns(const LLVMTypeConverter &converter, RewritePatternSet &patterns)
Populate the given list with patterns that convert from Math to ROCDL calls.
Lowering for gpu.dynamic.shared.memory to LLVM dialect.
The lowering of gpu.printf to a call to HIP hostcalls.
The lowering of gpu.printf to a call to an external printf() function.
Represents the amdgpu gfx chipset version, e.g., gfx90a, gfx942, gfx1103.
static FailureOr< Chipset > parse(StringRef name)
Parses the chipset version string and returns the chipset on success, and failure otherwise.