GPU_Dialect, CPred<"$_self.isa<::mlir::gpu::AsyncTokenType>()">, "async token type">,
BuildableType<"mlir::gpu::AsyncTokenType::get($_builder.getContext())">;
-/// Device-side synchronization token.
-def GPU_DeviceAsyncToken : DialectType<
- GPU_Dialect, CPred<"$_self.isa<::mlir::gpu::DeviceAsyncTokenType>()">,
- "device async token type">,
- BuildableType<
- "mlir::gpu::DeviceAsyncTokenType::get($_builder.getContext())">;
-
// Predicat to check if type is gpu::MMAMatrixType.
def IsMMAMatrixTypePred : CPred<"$_self.isa<::mlir::gpu::MMAMatrixType>()">;
using Base::Base;
};
-/// Device-side token storage type. There is only one type of device-side token.
-class DeviceAsyncTokenType
- : public Type::TypeBase<DeviceAsyncTokenType, Type, TypeStorage> {
-public:
- // Used for generic hooks in TypeBase.
- using Base::Base;
-};
-
/// MMAMatrixType storage and uniquing. Array is uniqued based on its shape
/// and type.
struct MMAMatrixStorageType : public TypeStorage {
}];
}
-def GPU_DeviceAsyncCopyOp : GPU_Op<"device_async_copy",
- [AttrSizedOperandSegments]> {
- let summary = "device-side asynchronous copy";
- let description = [{
- The `gpu.device_async_copy` op initiates an asynchronous copy operation of
- `$size` elements from source to the destination without blocking the thread.
- The destination has to be in shared memory.
-
- This is memory access will be pending to be added to a group.
-
- This op is meant to be used with `gpu.device_async_create_group` and
- `gpu.device_async_wait` to synchronize copies as explained in those ops
- descriptions.
-
- In order to do a copy and wait for the result we need the following
- combination:
- ```
- // copy 1.
- %cp1 = gpu.device_async_copy %A[%c0], %B[%c0], 4 :memref<16xf32> to memref<16xf32, 3>
- // copy 2.
- %cp2 = gpu.device_async_copy %C[%c0], %D[%c0], 4 : memref<16xf32> to memref<16xf32, 3>
- // group 1 contains copy 1 and copy 2.
- %token1 = gpu.device_async_create_group %cp1, %cp2
- // copy 3.
- %cp3 = gpu.device_async_copy %E[%c0], %F[%c0], 4 : memref<16xf32> to memref<16xf32, 3>
- // group 2 contains copy 3.
- %token2 = gpu.device_async_create_group %cp3
- // after the wait copy 1 and copy 2 are complete.
- gpu.device_async_wait %token1
- // after the wait copy 3 is complete.
- gpu.device_async_wait %token2
- ```
-
- Example:
-
- ```mlir
- %0 = gpu.device_async_copy %src[%c0, %c0], %dst[%c0, %c0, %c0], 4 :
- memref<4x5xf32> to memref<2x7x5xf32, 3>
- ```
- }];
- let results = (outs GPU_DeviceAsyncToken:$asyncToken);
- let arguments = (ins Arg<AnyMemRef, "", [MemWrite]>:$dst,
- Variadic<Index>:$dstIndices,
- Arg<AnyMemRef, "", [MemRead]>:$src,
- Variadic<Index>:$srcIndices,
- IndexAttr:$numElements);
- let assemblyFormat = [{
- $src `[` $srcIndices `]` `,` $dst `[` $dstIndices `]` `,` $numElements
- attr-dict `:` type($src) `to` type($dst)
- }];
- let hasVerifier = 1;
-}
-
-def GPU_DeviceAsyncCreateGroupOp : GPU_Op<"device_async_create_group", []> {
- let summary = "device side asynchronous create group operation";
- let description = [{
- The `gpu.device_async_create_group` op creates a group of memory accesses
- containing all the pending `device_async_copy` operations associated with
- argument tokens. Each token can only be part of one group.
-
- It returns a token that can be use to wait until the group fully completes.
-
- This is meant to be used with `gpu.device_async_wait` to synchronize copies
- as explained in those ops descriptions.
-
- Groups are executed in the order they are created.
-
- Example:
-
- ```mlir
- %0 = gpu.device_async_create_group
- ```
- }];
- let results = (outs GPU_DeviceAsyncToken:$asyncToken);
- let arguments = (ins Variadic<GPU_DeviceAsyncToken>:$inputTokens);
- let assemblyFormat = [{
- $inputTokens attr-dict
- }];
-}
-
-def GPU_DeviceAsyncWaitOp : GPU_Op<"device_async_wait", []> {
- let summary = "Wait for async gpu ops to complete.";
- let description = [{
- The `gpu.device_async_wait` op will block the execution thread until the group
- associated with the source token is fully completed.
-
- The optional `$numGroup` attribute gives a lower bound of the number of
- groups uncompleted when the wait can unblock the thread.
- Example:
-
- ```mlir
- gpu.device_async_wait %0
- ```
- }];
- let arguments = (ins GPU_DeviceAsyncToken:$asyncDependencies,
- OptionalAttr<I32Attr>:$numGroups);
- let assemblyFormat = [{
- $asyncDependencies attr-dict
- }];
-}
-
#endif // GPU_OPS
namespace mlir {
namespace NVVM {
+
+/// NVVM memory space identifiers.
+enum NVVMMemorySpace {
+ /// Global memory space identifier.
+ kGlobalMemorySpace = 1,
+ /// Shared memory space identifier.
+ kSharedMemorySpace = 3
+};
+
/// Return the element type and number of elements associated with a wmma matrix
/// of given chracteristics. This matches the logic in IntrinsicsNVVM.td
/// WMMA_REGS structure.
representing PTX specific operations while using MLIR high level concepts
like memref and 2-D vector.
}];
+ let useDefaultAttributePrinterParser = 1;
}
+/// Device-side synchronization token.
+def NVGPU_DeviceAsyncToken : DialectType<
+ NVGPU_Dialect, CPred<"$_self.isa<::mlir::nvgpu::DeviceAsyncTokenType>()">,
+ "device async token type">,
+ BuildableType<
+ "mlir::nvgpu::DeviceAsyncTokenType::get($_builder.getContext())">;
+
+
//===----------------------------------------------------------------------===//
// NVGPU Op definitions
//===----------------------------------------------------------------------===//
let description = [{
The `nvgpu.mma.sync` op represents the distributed form of a collective
matrix-multiply-and-accumulate (mma) operation that is compatible with
- `nvvm.mma.sync`. The operands and results are fragments of the full matrix
+ `nvvm.mma.sync`. The operands and results are fragments of the full matrix
operands. The full shape of the distributed mma operation is given by the
- `mmaShape` attribute in the form of a list of dimensions `[m, n, k]`.
+ `mmaShape` attribute in the form of a list of dimensions `[m, n, k]`.
This operation is meant to be lowered to the `nvvm.mma.sync` instruction, and
is an intermediate point between lowering from `vector.contract` to
`nvvm.mma.sync`.
-
+
This operation is meant to follow the semantic of described here:
https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma
-
+
Example:
-
+
```mlir
nvgpu.mma.sync (%a, %b, %c) :
(vector<4x2xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>
```
- }];
+ }];
let arguments = (ins AnyVector:$matrixA, AnyVector:$matrixB,
AnyVector:$matrixC, I64ArrayAttr:$mmaShape);
}];
}
+
+def NVGPU_DeviceAsyncCopyOp : NVGPU_Op<"device_async_copy",
+ [AttrSizedOperandSegments]> {
+ let summary = "device-side asynchronous copy";
+ let description = [{
+ The `gpu.device_async_copy` op initiates an asynchronous copy operation of
+ `$size` elements from source to the destination without blocking the thread.
+ The destination has to be in shared memory.
+
+ This is memory access will be pending to be added to a group.
+
+ This op is meant to be used with `gpu.device_async_create_group` and
+ `gpu.device_async_wait` to synchronize copies as explained in those ops
+ descriptions.
+ `bypassL1` attribute is hint to the backend and hardware that
+ the copy should by pass the L1 cache, this may be dropped by the backend or
+ hardware.
+
+ In order to do a copy and wait for the result we need the following
+ combination:
+ ```
+ // copy 1.
+ %cp1 = gpu.device_async_copy %A[%c0], %B[%c0], 4 :memref<16xf32> to memref<16xf32, 3>
+ // copy 2.
+ %cp2 = gpu.device_async_copy %C[%c0], %D[%c0], 4 : memref<16xf32> to memref<16xf32, 3>
+ // group 1 contains copy 1 and copy 2.
+ %token1 = gpu.device_async_create_group %cp1, %cp2
+ // copy 3.
+ %cp3 = gpu.device_async_copy %E[%c0], %F[%c0], 4 : memref<16xf32> to memref<16xf32, 3>
+ // group 2 contains copy 3.
+ %token2 = gpu.device_async_create_group %cp3
+ // after the wait copy 1 and copy 2 are complete.
+ gpu.device_async_wait %token1
+ // after the wait copy 3 is complete.
+ gpu.device_async_wait %token2
+ ```
+
+ Example:
+
+ ```mlir
+ %0 = gpu.device_async_copy %src[%c0, %c0], %dst[%c0, %c0, %c0], 4 :
+ memref<4x5xf32> to memref<2x7x5xf32, 3>
+ ```
+ }];
+ let results = (outs NVGPU_DeviceAsyncToken:$asyncToken);
+ let arguments = (ins Arg<AnyMemRef, "", [MemWrite]>:$dst,
+ Variadic<Index>:$dstIndices,
+ Arg<AnyMemRef, "", [MemRead]>:$src,
+ Variadic<Index>:$srcIndices,
+ IndexAttr:$numElements,
+ OptionalAttr<UnitAttr>:$bypassL1);
+ let assemblyFormat = [{
+ $src `[` $srcIndices `]` `,` $dst `[` $dstIndices `]` `,` $numElements
+ attr-dict `:` type($src) `to` type($dst)
+ }];
+ let hasVerifier = 1;
+}
+
+def NVGPU_DeviceAsyncCreateGroupOp : NVGPU_Op<"device_async_create_group", []> {
+ let summary = "device side asynchronous create group operation";
+ let description = [{
+ The `gpu.device_async_create_group` op creates a group of memory accesses
+ containing all the pending `device_async_copy` operations associated with
+ argument tokens. Each token can only be part of one group.
+
+ It returns a token that can be use to wait until the group fully completes.
+
+ This is meant to be used with `gpu.device_async_wait` to synchronize copies
+ as explained in those ops descriptions.
+
+ Groups are executed in the order they are created.
+
+ Example:
+
+ ```mlir
+ %0 = gpu.device_async_create_group
+ ```
+ }];
+ let results = (outs NVGPU_DeviceAsyncToken:$asyncToken);
+ let arguments = (ins Variadic<NVGPU_DeviceAsyncToken>:$inputTokens);
+ let assemblyFormat = [{
+ $inputTokens attr-dict
+ }];
+}
+
+def NVGPU_DeviceAsyncWaitOp : NVGPU_Op<"device_async_wait", []> {
+ let summary = "Wait for async gpu ops to complete.";
+ let description = [{
+ The `gpu.device_async_wait` op will block the execution thread until the group
+ associated with the source token is fully completed.
+
+ The optional `$numGroup` attribute gives a lower bound of the number of
+ groups uncompleted when the wait can unblock the thread.
+ Example:
+
+ ```mlir
+ gpu.device_async_wait %0
+ ```
+ }];
+ let arguments = (ins NVGPU_DeviceAsyncToken:$asyncDependencies,
+ OptionalAttr<I32Attr>:$numGroups);
+ let assemblyFormat = [{
+ $asyncDependencies attr-dict
+ }];
+}
+
#endif // NVGPU
#include "mlir/IR/OpDefinition.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
+namespace mlir {
+namespace nvgpu {
+
+/// Device-side token storage type. There is only one type of device-side token.
+class DeviceAsyncTokenType
+ : public Type::TypeBase<DeviceAsyncTokenType, Type, TypeStorage> {
+public:
+ // Used for generic hooks in TypeBase.
+ using Base::Base;
+};
+
+} // namespace nvgpu
+} // namespace mlir
+
#include "mlir/Dialect/NVGPU/NVGPUDialect.h.inc"
#define GET_OP_CLASSES
namespace {
-/// NVVM memory space identifiers.
-enum NVVMMemorySpace {
- /// Global memory space identifier.
- kGlobalMemorySpace = 1,
- /// Shared memory space identifier.
- kSharedMemorySpace = 3
-};
-
/// Convert gpu dialect shfl mode enum to the equivalent nvvm one.
static NVVM::ShflKind convertShflKind(gpu::ShuffleMode mode) {
switch (mode) {
}
};
-struct GPUAsyncCopyLowering
- : public ConvertOpToLLVMPattern<gpu::DeviceAsyncCopyOp> {
- using ConvertOpToLLVMPattern<gpu::DeviceAsyncCopyOp>::ConvertOpToLLVMPattern;
-
- LogicalResult
- matchAndRewrite(gpu::DeviceAsyncCopyOp op, OpAdaptor adaptor,
- ConversionPatternRewriter &rewriter) const override {
- Location loc = op->getLoc();
- auto dstMemrefType = op.dst().getType().cast<MemRefType>();
- Value dstPtr = getStridedElementPtr(loc, dstMemrefType, adaptor.dst(),
- adaptor.dstIndices(), rewriter);
- auto i8Ty = IntegerType::get(op.getContext(), 8);
- auto dstPointerType =
- LLVM::LLVMPointerType::get(i8Ty, dstMemrefType.getMemorySpaceAsInt());
- dstPtr = rewriter.create<LLVM::BitcastOp>(loc, dstPointerType, dstPtr);
-
- auto srcMemrefType = op.src().getType().cast<MemRefType>();
-
- Value scrPtr = getStridedElementPtr(loc, srcMemrefType, adaptor.src(),
- adaptor.srcIndices(), rewriter);
- auto srcPointerType =
- LLVM::LLVMPointerType::get(i8Ty, srcMemrefType.getMemorySpaceAsInt());
- scrPtr = rewriter.create<LLVM::BitcastOp>(loc, srcPointerType, scrPtr);
- // Intrinsics takes a global pointer so we need an address space cast.
- auto srcPointerGlobalType =
- LLVM::LLVMPointerType::get(i8Ty, NVVMMemorySpace::kGlobalMemorySpace);
- scrPtr = rewriter.create<LLVM::AddrSpaceCastOp>(loc, srcPointerGlobalType,
- scrPtr);
- int64_t numElements = adaptor.numElements().getZExtValue();
- int64_t sizeInBytes =
- (dstMemrefType.getElementTypeBitWidth() / 8) * numElements;
- rewriter.create<NVVM::CpAsyncOp>(loc, dstPtr, scrPtr,
- rewriter.getI32IntegerAttr(sizeInBytes),
- /*bypassL1=*/UnitAttr());
-
- // Drop the result token.
- Value zero = rewriter.create<LLVM::ConstantOp>(
- op->getLoc(), IntegerType::get(op.getContext(), 32),
- rewriter.getI32IntegerAttr(0));
- rewriter.replaceOp(op, zero);
- return success();
- }
-};
-
-struct GPUAsyncCreateGroupLowering
- : public ConvertOpToLLVMPattern<gpu::DeviceAsyncCreateGroupOp> {
- using ConvertOpToLLVMPattern<
- gpu::DeviceAsyncCreateGroupOp>::ConvertOpToLLVMPattern;
-
- LogicalResult
- matchAndRewrite(gpu::DeviceAsyncCreateGroupOp op, OpAdaptor adaptor,
- ConversionPatternRewriter &rewriter) const override {
- rewriter.create<NVVM::CpAsyncCommitGroupOp>(op.getLoc());
- // Drop the result token.
- Value zero = rewriter.create<LLVM::ConstantOp>(
- op->getLoc(), IntegerType::get(op.getContext(), 32),
- rewriter.getI32IntegerAttr(0));
- rewriter.replaceOp(op, zero);
- return success();
- }
-};
-
-struct GPUAsyncWaitLowering
- : public ConvertOpToLLVMPattern<gpu::DeviceAsyncWaitOp> {
- using ConvertOpToLLVMPattern<gpu::DeviceAsyncWaitOp>::ConvertOpToLLVMPattern;
-
- LogicalResult
- matchAndRewrite(gpu::DeviceAsyncWaitOp op, OpAdaptor adaptor,
- ConversionPatternRewriter &rewriter) const override {
- // If numGroup is not present pick 0 as a conservative correct value.
- int32_t numGroups = adaptor.numGroups() ? *adaptor.numGroups() : 0;
- rewriter.create<NVVM::CpAsyncWaitGroupOp>(op.getLoc(), numGroups);
- rewriter.eraseOp(op);
- return success();
- }
-};
-
struct GPULaneIdOpToNVVM : ConvertOpToLLVMPattern<gpu::LaneIdOp> {
using ConvertOpToLLVMPattern<gpu::LaneIdOp>::ConvertOpToLLVMPattern;
return llvm::None;
return converter.convertType(MemRefType::Builder(type).setMemorySpace(0));
});
- /// device-side async tokens cannot be materialized in nvvm. We just convert
- /// them to a dummy i32 type in order to easily drop them during conversion.
- converter.addConversion([&](gpu::DeviceAsyncTokenType type) -> Type {
- return converter.convertType(IntegerType::get(type.getContext(), 32));
- });
// Lowering for MMAMatrixType.
converter.addConversion([&](gpu::MMAMatrixType type) -> Type {
return convertMMAToLLVMType(type);
"__nv_sqrt");
patterns.add<OpToFuncCallLowering<math::TanhOp>>(converter, "__nv_tanhf",
"__nv_tanh");
- patterns.add<GPUAsyncCopyLowering, GPUAsyncCreateGroupLowering,
- GPUAsyncWaitLowering>(converter);
}
std::unique_ptr<OperationPass<gpu::GPUModuleOp>>
Core
LINK_LIBS PUBLIC
+ MLIRGPUOps
MLIRLLVMCommonConversion
MLIRLLVMIR
MLIRNVVMIR
#include "../PassDetail.h"
#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Dialect/GPU/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
#include "mlir/Dialect/NVGPU/NVGPUDialect.h"
void runOnOperation() override {
RewritePatternSet patterns(&getContext());
LLVMTypeConverter converter(&getContext());
+ /// device-side async tokens cannot be materialized in nvvm. We just convert
+ /// them to a dummy i32 type in order to easily drop them during conversion.
+ converter.addConversion([&](nvgpu::DeviceAsyncTokenType type) -> Type {
+ return converter.convertType(IntegerType::get(type.getContext(), 32));
+ });
populateNVGPUToNVVMConversionPatterns(converter, patterns);
LLVMConversionTarget target(getContext());
target.addLegalDialect<::mlir::LLVM::LLVMDialect>();
}
};
+struct NVGPUAsyncCopyLowering
+ : public ConvertOpToLLVMPattern<nvgpu::DeviceAsyncCopyOp> {
+ using ConvertOpToLLVMPattern<
+ nvgpu::DeviceAsyncCopyOp>::ConvertOpToLLVMPattern;
+
+ LogicalResult
+ matchAndRewrite(nvgpu::DeviceAsyncCopyOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ Location loc = op->getLoc();
+ auto dstMemrefType = op.dst().getType().cast<MemRefType>();
+ Value dstPtr = getStridedElementPtr(loc, dstMemrefType, adaptor.dst(),
+ adaptor.dstIndices(), rewriter);
+ auto i8Ty = IntegerType::get(op.getContext(), 8);
+ auto dstPointerType =
+ LLVM::LLVMPointerType::get(i8Ty, dstMemrefType.getMemorySpaceAsInt());
+ dstPtr = rewriter.create<LLVM::BitcastOp>(loc, dstPointerType, dstPtr);
+
+ auto srcMemrefType = op.src().getType().cast<MemRefType>();
+
+ Value scrPtr = getStridedElementPtr(loc, srcMemrefType, adaptor.src(),
+ adaptor.srcIndices(), rewriter);
+ auto srcPointerType =
+ LLVM::LLVMPointerType::get(i8Ty, srcMemrefType.getMemorySpaceAsInt());
+ scrPtr = rewriter.create<LLVM::BitcastOp>(loc, srcPointerType, scrPtr);
+ // Intrinsics takes a global pointer so we need an address space cast.
+ auto srcPointerGlobalType = LLVM::LLVMPointerType::get(
+ i8Ty, NVVM::NVVMMemorySpace::kGlobalMemorySpace);
+ scrPtr = rewriter.create<LLVM::AddrSpaceCastOp>(loc, srcPointerGlobalType,
+ scrPtr);
+ int64_t numElements = adaptor.numElements().getZExtValue();
+ int64_t sizeInBytes =
+ (dstMemrefType.getElementTypeBitWidth() / 8) * numElements;
+ // bypass L1 is only supported for byte sizes of 16, we drop the hint
+ // otherwise.
+ UnitAttr bypassL1 = sizeInBytes == 16 ? adaptor.bypassL1Attr() : UnitAttr();
+ rewriter.create<NVVM::CpAsyncOp>(
+ loc, dstPtr, scrPtr, rewriter.getI32IntegerAttr(sizeInBytes), bypassL1);
+
+ // Drop the result token.
+ Value zero = rewriter.create<LLVM::ConstantOp>(
+ op->getLoc(), IntegerType::get(op.getContext(), 32),
+ rewriter.getI32IntegerAttr(0));
+ rewriter.replaceOp(op, zero);
+ return success();
+ }
+};
+
+struct NVGPUAsyncCreateGroupLowering
+ : public ConvertOpToLLVMPattern<nvgpu::DeviceAsyncCreateGroupOp> {
+ using ConvertOpToLLVMPattern<
+ nvgpu::DeviceAsyncCreateGroupOp>::ConvertOpToLLVMPattern;
+
+ LogicalResult
+ matchAndRewrite(nvgpu::DeviceAsyncCreateGroupOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ rewriter.create<NVVM::CpAsyncCommitGroupOp>(op.getLoc());
+ // Drop the result token.
+ Value zero = rewriter.create<LLVM::ConstantOp>(
+ op->getLoc(), IntegerType::get(op.getContext(), 32),
+ rewriter.getI32IntegerAttr(0));
+ rewriter.replaceOp(op, zero);
+ return success();
+ }
+};
+
+struct NVGPUAsyncWaitLowering
+ : public ConvertOpToLLVMPattern<nvgpu::DeviceAsyncWaitOp> {
+ using ConvertOpToLLVMPattern<
+ nvgpu::DeviceAsyncWaitOp>::ConvertOpToLLVMPattern;
+
+ LogicalResult
+ matchAndRewrite(nvgpu::DeviceAsyncWaitOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ // If numGroup is not present pick 0 as a conservative correct value.
+ int32_t numGroups = adaptor.numGroups() ? *adaptor.numGroups() : 0;
+ rewriter.create<NVVM::CpAsyncWaitGroupOp>(op.getLoc(), numGroups);
+ rewriter.eraseOp(op);
+ return success();
+ }
+};
+
} // namespace
+
void mlir::populateNVGPUToNVVMConversionPatterns(LLVMTypeConverter &converter,
RewritePatternSet &patterns) {
- patterns.add<MmaSyncOptoNVVM, MmaLdMatrixOpToNVVM>(converter);
+ patterns.add<MmaSyncOptoNVVM, MmaLdMatrixOpToNVVM, NVGPUAsyncCopyLowering,
+ NVGPUAsyncCreateGroupLowering, NVGPUAsyncWaitLowering>(
+ converter);
}
std::unique_ptr<Pass> mlir::createConvertNVGPUToNVVMPass() {
void GPUDialect::initialize() {
addTypes<AsyncTokenType>();
- addTypes<DeviceAsyncTokenType>();
addTypes<MMAMatrixType>();
addOperations<
#define GET_OP_LIST
// Handle 'async token' types.
if (keyword == "async.token")
return AsyncTokenType::get(context);
- // Handle 'device async token' types.
- if (keyword == "device.async.token")
- return DeviceAsyncTokenType::get(context);
if (keyword == "mma_matrix") {
SMLoc beginLoc = parser.getNameLoc();
void GPUDialect::printType(Type type, DialectAsmPrinter &os) const {
TypeSwitch<Type>(type)
.Case<AsyncTokenType>([&](Type) { os << "async.token"; })
- .Case<DeviceAsyncTokenType>([&](Type) { os << "device.async.token"; })
.Case<MMAMatrixType>([&](MMAMatrixType fragTy) {
os << "mma_matrix<";
auto shape = fragTy.getShape();
results.add<SimplifyDimOfAllocOp>(context);
}
-//===----------------------------------------------------------------------===//
-// GPU_DeviceAsyncCopyOp
-//===----------------------------------------------------------------------===//
-
-LogicalResult DeviceAsyncCopyOp::verify() {
- auto srcMemref = src().getType().cast<MemRefType>();
- auto dstMemref = dst().getType().cast<MemRefType>();
- unsigned workgroupAddressSpace = GPUDialect::getWorkgroupAddressSpace();
- if (!isLastMemrefDimUnitStride(srcMemref))
- return emitError("source memref most minor dim must have unit stride");
- if (!isLastMemrefDimUnitStride(dstMemref))
- return emitError("destination memref most minor dim must have unit stride");
- if (dstMemref.getMemorySpaceAsInt() != workgroupAddressSpace)
- return emitError("destination memref must have memory space ")
- << workgroupAddressSpace;
- if (dstMemref.getElementType() != srcMemref.getElementType())
- return emitError("source and destination must have the same element type");
- if (size_t(srcMemref.getRank()) != srcIndices().size())
- return emitOpError() << "expected " << srcMemref.getRank()
- << " source indices, got " << srcIndices().size();
- if (size_t(dstMemref.getRank()) != dstIndices().size())
- return emitOpError() << "expected " << dstMemref.getRank()
- << " destination indices, got " << dstIndices().size();
- return success();
-}
-
#include "mlir/Dialect/GPU/GPUOpInterfaces.cpp.inc"
#include "mlir/Dialect/GPU/GPUOpsEnums.cpp.inc"
MLIRNVGPUIncGen
LINK_LIBS PUBLIC
+ MLIRGPUOps
MLIRIR
MLIRSideEffectInterfaces
)
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/NVGPU/NVGPUDialect.h"
+#include "mlir/Dialect/GPU/GPUDialect.h"
#include "mlir/IR/Builders.h"
+#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/TypeUtilities.h"
+#include "llvm/ADT/TypeSwitch.h"
using namespace mlir;
+using namespace mlir::nvgpu;
#include "mlir/Dialect/NVGPU/NVGPUDialect.cpp.inc"
void nvgpu::NVGPUDialect::initialize() {
+ addTypes<DeviceAsyncTokenType>();
addOperations<
#define GET_OP_LIST
#include "mlir/Dialect/NVGPU/NVGPU.cpp.inc"
>();
}
+Type NVGPUDialect::parseType(DialectAsmParser &parser) const {
+ // Parse the main keyword for the type.
+ StringRef keyword;
+ if (parser.parseKeyword(&keyword))
+ return Type();
+ MLIRContext *context = getContext();
+ // Handle 'device async token' types.
+ if (keyword == "device.async.token")
+ return DeviceAsyncTokenType::get(context);
+
+ parser.emitError(parser.getNameLoc(), "unknown nvgpu type: " + keyword);
+ return Type();
+}
+
+void NVGPUDialect::printType(Type type, DialectAsmPrinter &os) const {
+ TypeSwitch<Type>(type)
+ .Case<DeviceAsyncTokenType>([&](Type) { os << "device.async.token"; })
+ .Default([](Type) { llvm_unreachable("unexpected 'nvgpu' type kind"); });
+}
+//===----------------------------------------------------------------------===//
+// NVGPU_DeviceAsyncCopyOp
+//===----------------------------------------------------------------------===//
+
+/// Return true if the last dimension of the MemRefType has unit stride. Also
+/// return true for memrefs with no strides.
+static bool isLastMemrefDimUnitStride(MemRefType type) {
+ int64_t offset;
+ SmallVector<int64_t> strides;
+ if (failed(getStridesAndOffset(type, strides, offset))) {
+ return false;
+ }
+ return strides.back() == 1;
+}
+
+LogicalResult DeviceAsyncCopyOp::verify() {
+ auto srcMemref = src().getType().cast<MemRefType>();
+ auto dstMemref = dst().getType().cast<MemRefType>();
+ unsigned workgroupAddressSpace = gpu::GPUDialect::getWorkgroupAddressSpace();
+ if (!isLastMemrefDimUnitStride(srcMemref))
+ return emitError("source memref most minor dim must have unit stride");
+ if (!isLastMemrefDimUnitStride(dstMemref))
+ return emitError("destination memref most minor dim must have unit stride");
+ if (dstMemref.getMemorySpaceAsInt() != workgroupAddressSpace)
+ return emitError("destination memref must have memory space ")
+ << workgroupAddressSpace;
+ if (dstMemref.getElementType() != srcMemref.getElementType())
+ return emitError("source and destination must have the same element type");
+ if (size_t(srcMemref.getRank()) != srcIndices().size())
+ return emitOpError() << "expected " << srcMemref.getRank()
+ << " source indices, got " << srcIndices().size();
+ if (size_t(dstMemref.getRank()) != dstIndices().size())
+ return emitOpError() << "expected " << dstMemref.getRank()
+ << " destination indices, got " << dstIndices().size();
+ return success();
+}
+
#define GET_OP_CLASSES
#include "mlir/Dialect/NVGPU/NVGPU.cpp.inc"
}
}
-// -----
-
-gpu.module @test_module {
- // CHECK-LABEL: @async_cp(
- // CHECK-SAME: %[[IDX:[a-zA-Z0-9_]+]]: i64)
- gpu.func @async_cp(
- %src: memref<128x128xf32>, %dst: memref<3x16x128xf32, 3>, %i : index) kernel {
- // CHECK-DAG: %[[BASEDST:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32, 3>, ptr<f32, 3>, i64, array<3 x i64>, array<3 x i64>)>
- // CHECK-DAG: %[[S0:.*]] = llvm.mlir.constant(2048 : index) : i64
- // CHECK-DAG: %[[LI:.*]] = llvm.mul %[[IDX]], %[[S0]] : i64
- // CHECK-DAG: %[[S1:.*]] = llvm.mlir.constant(128 : index) : i64
- // CHECK-DAG: %[[FI0:.*]] = llvm.mul %[[IDX]], %[[S1]] : i64
- // CHECK-DAG: %[[FI1:.*]] = llvm.add %[[LI]], %[[FI0]] : i64
- // CHECK-DAG: %[[FI2:.*]] = llvm.add %[[FI1]], %[[IDX]] : i64
- // CHECK-DAG: %[[ADDRESSDST:.*]] = llvm.getelementptr %[[BASEDST]][%[[FI2]]] : (!llvm.ptr<f32, 3>, i64) -> !llvm.ptr<f32, 3>
- // CHECK-DAG: %[[CAST0:.*]] = llvm.bitcast %[[ADDRESSDST]] : !llvm.ptr<f32, 3> to !llvm.ptr<i8, 3>
- // CHECK-DAG: %[[BASESRC:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
- // CHECK-DAG: %[[S3:.*]] = llvm.mlir.constant(128 : index) : i64
- // CHECK-DAG: %[[FI3:.*]] = llvm.mul %[[IDX]], %[[S3]] : i64
- // CHECK-DAG: %[[FI4:.*]] = llvm.add %[[FI3]], %[[IDX]] : i64
- // CHECK-DAG: %[[ADDRESSSRC:.*]] = llvm.getelementptr %[[BASESRC]][%[[FI4]]] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
- // CHECK-DAG: %[[CAST1:.*]] = llvm.bitcast %[[ADDRESSSRC]] : !llvm.ptr<f32> to !llvm.ptr<i8>
- // CHECK-DAG: %[[CAST2:.*]] = llvm.addrspacecast %[[CAST1]] : !llvm.ptr<i8> to !llvm.ptr<i8, 1>
- // CHECK-DAG: nvvm.cp.async.shared.global %[[CAST0]], %[[CAST2]], 16
- %0 = gpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 4 : memref<128x128xf32> to memref<3x16x128xf32, 3>
- // CHECK: nvvm.cp.async.commit.group
- %1 = gpu.device_async_create_group %0
- // CHECK: nvvm.cp.async.wait.group 1
- gpu.device_async_wait %1 { numGroups = 1 : i32 }
- gpu.return
- }
-}
// CHECK: llvm.bitcast [[el]] : f32 to vector<1xf32>
// CHECK-COUNT-4: llvm.insertvalue {{.*}} : !llvm.array<4 x vector<1xf32>>
return %d : vector<4x1xf32>
-}
\ No newline at end of file
+}
+
+// -----
+
+// CHECK-LABEL: @async_cp(
+// CHECK-SAME: %[[IDX:[a-zA-Z0-9_]+]]: index)
+func.func @async_cp(
+ %src: memref<128x128xf32>, %dst: memref<3x16x128xf32, 3>, %i : index) {
+ // CHECK: %[[IDX1:.*]] = builtin.unrealized_conversion_cast %[[IDX]] : index to i64
+ // CHECK-DAG: %[[BASEDST:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32, 3>, ptr<f32, 3>, i64, array<3 x i64>, array<3 x i64>)>
+ // CHECK-DAG: %[[S0:.*]] = llvm.mlir.constant(2048 : index) : i64
+ // CHECK-DAG: %[[LI:.*]] = llvm.mul %[[IDX1]], %[[S0]] : i64
+ // CHECK-DAG: %[[S1:.*]] = llvm.mlir.constant(128 : index) : i64
+ // CHECK-DAG: %[[FI0:.*]] = llvm.mul %[[IDX1]], %[[S1]] : i64
+ // CHECK-DAG: %[[FI1:.*]] = llvm.add %[[LI]], %[[FI0]] : i64
+ // CHECK-DAG: %[[FI2:.*]] = llvm.add %[[FI1]], %[[IDX1]] : i64
+ // CHECK-DAG: %[[ADDRESSDST:.*]] = llvm.getelementptr %[[BASEDST]][%[[FI2]]] : (!llvm.ptr<f32, 3>, i64) -> !llvm.ptr<f32, 3>
+ // CHECK-DAG: %[[CAST0:.*]] = llvm.bitcast %[[ADDRESSDST]] : !llvm.ptr<f32, 3> to !llvm.ptr<i8, 3>
+ // CHECK-DAG: %[[BASESRC:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+ // CHECK-DAG: %[[S3:.*]] = llvm.mlir.constant(128 : index) : i64
+ // CHECK-DAG: %[[FI3:.*]] = llvm.mul %[[IDX1]], %[[S3]] : i64
+ // CHECK-DAG: %[[FI4:.*]] = llvm.add %[[FI3]], %[[IDX1]] : i64
+ // CHECK-DAG: %[[ADDRESSSRC:.*]] = llvm.getelementptr %[[BASESRC]][%[[FI4]]] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
+ // CHECK-DAG: %[[CAST1:.*]] = llvm.bitcast %[[ADDRESSSRC]] : !llvm.ptr<f32> to !llvm.ptr<i8>
+ // CHECK-DAG: %[[CAST2:.*]] = llvm.addrspacecast %[[CAST1]] : !llvm.ptr<i8> to !llvm.ptr<i8, 1>
+ // CHECK-DAG: nvvm.cp.async.shared.global %[[CAST0]], %[[CAST2]], 16
+ %0 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 4 : memref<128x128xf32> to memref<3x16x128xf32, 3>
+ // CHECK: nvvm.cp.async.commit.group
+ %1 = nvgpu.device_async_create_group %0
+ // CHECK: nvvm.cp.async.wait.group 1
+ nvgpu.device_async_wait %1 { numGroups = 1 : i32 }
+
+ // CHECK: nvvm.cp.async.shared.global %{{.*}}, %{{.*}}, 16 {bypass_l1}
+ %2 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 4 {bypassL1}: memref<128x128xf32> to memref<3x16x128xf32, 3>
+ return
+}
+
// -----
-func.func @async_cp_memory_space(%dst : memref<16xf32>, %src : memref<16xf32>, %i : index) -> () {
- // expected-error @+1 {{destination memref must have memory space 3}}
- gpu.device_async_copy %src[%i], %dst[%i], 16 : memref<16xf32> to memref<16xf32>
- return
-}
-
-// -----
-
-func.func @async_cp_memref_type(%dst : memref<16xi32, 3>, %src : memref<16xf32>, %i : index) -> () {
- // expected-error @+1 {{source and destination must have the same element type}}
- gpu.device_async_copy %src[%i], %dst[%i], 16 : memref<16xf32> to memref<16xi32, 3>
- return
-}
-
-// -----
-
-func.func @async_cp_num_src_indices(%dst : memref<16xf32, 3>, %src : memref<16x16xf32>, %i : index) -> () {
- // expected-error @+1 {{expected 2 source indices, got 1}}
- gpu.device_async_copy %src[%i], %dst[%i], 16 : memref<16x16xf32> to memref<16xf32, 3>
- return
-}
-
-// -----
-
-func.func @async_cp_num_dst_indices(%dst : memref<16x16xf32, 3>, %src : memref<16xf32>, %i : index) -> () {
- // expected-error @+1 {{expected 2 destination indices, got 1}}
- gpu.device_async_copy %src[%i], %dst[%i], 16 : memref<16xf32> to memref<16x16xf32, 3>
- return
-}
-
-// -----
-
-func.func @async_cp_num_src_stride(
- %dst : memref<200x100xf32, 3>,
- %src : memref<200x100xf32, affine_map<(d0, d1) -> (200*d0 + 2*d1)>>,
- %i : index) -> () {
- // expected-error @+1 {{source memref most minor dim must have unit stride}}
- gpu.device_async_copy %src[%i, %i], %dst[%i, %i], 16 :
- memref<200x100xf32, affine_map<(d0, d1) -> (200*d0 + 2*d1)>> to memref<200x100xf32, 3>
- return
-}
-
-// -----
-
-func.func @async_cp_num_dst_stride(
- %dst : memref<200x100xf32, affine_map<(d0, d1) -> (200*d0 + 2*d1)>, 3>,
- %src : memref<200x100xf32>,
- %i : index) -> () {
- // expected-error @+1 {{destination memref most minor dim must have unit stride}}
- gpu.device_async_copy %src[%i, %i], %dst[%i, %i], 16 :
- memref<200x100xf32> to memref<200x100xf32, affine_map<(d0, d1) -> (200*d0 + 2*d1)>, 3>
- return
-}
-
-// -----
-
// Number of symbol operand count less than memref symbol count.
func.func @alloc() {
// expected-error@+1 {{symbol operand count does not equal memref symbol count}}
return
}
- func.func @async_cp(%dst : memref<2x7x5xf32, 3>, %src : memref<4x5xf32>){
- // CHECK-LABEL: func @async_cp
- %c0 = arith.constant 0 : index
- // CHECK: gpu.device_async_copy %{{.*}}[{{.*}}, {{.*}}], %{{.*}}[{{.*}}, {{.*}}, {{.*}}], 4 : memref<4x5xf32> to memref<2x7x5xf32, 3>
- %0 = gpu.device_async_copy %src[%c0, %c0], %dst[%c0, %c0, %c0], 4 : memref<4x5xf32> to memref<2x7x5xf32, 3>
- // CHECK: %{{.*}} = gpu.device_async_create_group
- %token = gpu.device_async_create_group %0
- // CHECK: gpu.device_async_wait %{{.*}} {numGroups = 1 : i32}
- gpu.device_async_wait %token {numGroups = 1 : i32}
- return
- }
-
// CHECK-LABEL: func @set_default_device
func.func @set_default_device(%arg0: i32) {
// CHECK: gpu.set_default_device
--- /dev/null
+// RUN: mlir-opt -split-input-file -verify-diagnostics %s
+
+func.func @async_cp_memory_space(%dst : memref<16xf32>, %src : memref<16xf32>, %i : index) -> () {
+ // expected-error @+1 {{destination memref must have memory space 3}}
+ nvgpu.device_async_copy %src[%i], %dst[%i], 16 : memref<16xf32> to memref<16xf32>
+ return
+}
+
+// -----
+
+func.func @async_cp_memref_type(%dst : memref<16xi32, 3>, %src : memref<16xf32>, %i : index) -> () {
+ // expected-error @+1 {{source and destination must have the same element type}}
+ nvgpu.device_async_copy %src[%i], %dst[%i], 16 : memref<16xf32> to memref<16xi32, 3>
+ return
+}
+
+// -----
+
+func.func @async_cp_num_src_indices(%dst : memref<16xf32, 3>, %src : memref<16x16xf32>, %i : index) -> () {
+ // expected-error @+1 {{expected 2 source indices, got 1}}
+ nvgpu.device_async_copy %src[%i], %dst[%i], 16 : memref<16x16xf32> to memref<16xf32, 3>
+ return
+}
+
+// -----
+
+func.func @async_cp_num_dst_indices(%dst : memref<16x16xf32, 3>, %src : memref<16xf32>, %i : index) -> () {
+ // expected-error @+1 {{expected 2 destination indices, got 1}}
+ nvgpu.device_async_copy %src[%i], %dst[%i], 16 : memref<16xf32> to memref<16x16xf32, 3>
+ return
+}
+
+// -----
+
+func.func @async_cp_num_src_stride(
+ %dst : memref<200x100xf32, 3>,
+ %src : memref<200x100xf32, affine_map<(d0, d1) -> (200*d0 + 2*d1)>>,
+ %i : index) -> () {
+ // expected-error @+1 {{source memref most minor dim must have unit stride}}
+ nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i], 16 :
+ memref<200x100xf32, affine_map<(d0, d1) -> (200*d0 + 2*d1)>> to memref<200x100xf32, 3>
+ return
+}
+
+// -----
+
+func.func @async_cp_num_dst_stride(
+ %dst : memref<200x100xf32, affine_map<(d0, d1) -> (200*d0 + 2*d1)>, 3>,
+ %src : memref<200x100xf32>,
+ %i : index) -> () {
+ // expected-error @+1 {{destination memref most minor dim must have unit stride}}
+ nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i], 16 :
+ memref<200x100xf32> to memref<200x100xf32, affine_map<(d0, d1) -> (200*d0 + 2*d1)>, 3>
+ return
+}
func.func @mma_sync(%arg0: vector<4x2xf16>,
%arg1: vector<2x2xf16>,
%arg2: vector<2x2xf16>) -> vector<2x2xf16> {
-// CHECK: nvgpu.mma.sync(%{{.*}}, %{{.*}}, %{{.*}}) {mmaShape = [16, 8, 16]} : (vector<4x2xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>
+// CHECK: nvgpu.mma.sync(%{{.*}}, %{{.*}}, %{{.*}}) {mmaShape = [16, 8, 16]} : (vector<4x2xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>
%d = nvgpu.mma.sync(%arg0, %arg1, %arg2) {mmaShape = [16, 8, 16]} :
- (vector<4x2xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>
+ (vector<4x2xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>
return %d : vector<2x2xf16>
}
+
+
+func.func @async_cp(%dst : memref<2x7x5xf32, 3>, %src : memref<4x5xf32>){
+ // CHECK-LABEL: func @async_cp
+ %c0 = arith.constant 0 : index
+ // CHECK: nvgpu.device_async_copy %{{.*}}[{{.*}}, {{.*}}], %{{.*}}[{{.*}}, {{.*}}, {{.*}}], 4 : memref<4x5xf32> to memref<2x7x5xf32, 3>
+ %0 = nvgpu.device_async_copy %src[%c0, %c0], %dst[%c0, %c0, %c0], 4 : memref<4x5xf32> to memref<2x7x5xf32, 3>
+ // CHECK: %{{.*}} = nvgpu.device_async_create_group
+ %token = nvgpu.device_async_create_group %0
+ // CHECK: nvgpu.device_async_wait %{{.*}} {numGroups = 1 : i32}
+ nvgpu.device_async_wait %token {numGroups = 1 : i32}
+ return
+}
hdrs = ["include/mlir/Dialect/NVGPU/NVGPUDialect.h"],
includes = ["include"],
deps = [
+ ":GPUDialect",
":IR",
":NVGPUIncGen",
":SideEffectInterfaces",
includes = ["include"],
deps = [
":ConversionPassIncGen",
+ ":GPUDialect",
":IR",
":LLVMCommonConversion",
":NVGPU",
projects / platform / upstream / llvm.git / commitdiff