Page MenuHomePhabricator
Differential D95331 Diff 345358 mlir/lib/Conversion/GPUToNVVM/WmmaOpsToNvvmLowering.cpp

Changeset View

Standalone View

View Options

mlir/lib/Conversion/GPUToNVVM/WmmaOpsToNvvmLowering.cpp

  • This file was added.
//===-- WmmaOpsToNVVMLowering.cpp - WMMA LD/ST/Compute to NVVM 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 definitions of patterns to lower GPU Subgroup MMA ops to
// NVVM Dialect.
//
//===----------------------------------------------------------------------===//
#include "CommonTypes.h"
Lint: Pre-merge checks
Inline Actions

clang-format: please reformat the code

-#include "CommonTypes.h"
Lint: Pre-merge checks: clang-format: please reformat the code ``` -#include "CommonTypes.h" ```
#include "WmmaOpsToNvvmLowering.h"
Lint: Pre-merge checks
Inline Actions

clang-format: please reformat the code

+#include "CommonTypes.h"
Lint: Pre-merge checks: clang-format: please reformat the code ``` +#include "CommonTypes.h" ```
#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
#include "mlir/Dialect/GPU/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
namespace mlir {
/// Checks if all the operands of the op being lowered are of LLVM Types. The
/// types are expected to be converted by the `LLVMTypeConverter` before the op
/// is actually lowered. If the type of an operands is not already converted it
/// hints a missing typeConversion and failure is returned in that case.
static LogicalResult areAllLLVMTypes(Operation *op, ValueRange operands,
ConversionPatternRewriter &rewriter) {
if (!llvm::all_of(operands, [](Value value) {
return LLVM::isCompatibleType(value.getType());
}))
return rewriter.notifyMatchFailure(
ftynseUnsubmitted
Done
ReplyInline Actions

Nit: add braces here, this isn't a trivial condition anymore.

ftynse: Nit: add braces here, this isn't a trivial condition anymore.
op, "cannot convert if operands aren't of LLVM type.");
return success();
}
WmmaLoadOpToNVVMLowering::WmmaLoadOpToNVVMLowering(
LLVMTypeConverter &typeConverter)
: ConvertOpToLLVMPattern<gpu::SubgroupMmaLoadMatrixOp>(typeConverter),
CommonLLVMAndBuiltInMLIRTypes(&this->getTypeConverter()->getContext()) {}
LogicalResult WmmaLoadOpToNVVMLowering::matchAndRewrite(
gpu::SubgroupMmaLoadMatrixOp subgroupMmaLoadMatrixOp,
ArrayRef<Value> operands, ConversionPatternRewriter &rewriter) const {
Operation *op = subgroupMmaLoadMatrixOp.getOperation();
if (failed(areAllLLVMTypes(op, operands, rewriter)))
return failure();
int8_t indexTypeBitwidth = this->getTypeConverter()->getIndexTypeBitwidth();
ftynseUnsubmitted
Done
ReplyInline Actions

I would have just used unsigned here.

ftynse: I would have just used `unsigned` here.
// The corresponding intrinsics expects leadDimension to be a 32-bit
// integer, so all the calculations of linearizing the load address
// must also follow this restriction.
if (indexTypeBitwidth != 32)
return rewriter.notifyMatchFailure(
op, "Expected indices to the memref to be 32-bit wide.");
// Source memref of the original op.
MemRefType srcMemrefType =
subgroupMmaLoadMatrixOp.srcMemref().getType().cast<MemRefType>();
Location loc = op->getLoc();
auto leadDimension = subgroupMmaLoadMatrixOp.leadDimensionAttr();
// MemRefDescriptor to extract alignedPtr and offset.
MemRefDescriptor promotedSrcOp(
gpu::SubgroupMmaLoadMatrixOpAdaptor(operands).srcMemref());
// Emit ops which compute the load offset using `srcOffsetI`,
// `srcOffsetJ`. The actualOffset is (memrefOffset + (alignedPtr +
// ((leadDimension * srcOffsetI) + srcOffsetJ)). The memrefs here are
// assumed to be normalized and hence the simple conversion works.
SmallVector<Value> indices(subgroupMmaLoadMatrixOp.indices());
Value srcOffsetIVal = indices[0];
Value srcOffsetJVal = indices[1];
Value leadingDim32 =
rewriter.create<LLVM::ConstantOp>(loc, i32Ty, leadDimension);
Value numElemsLeadDim =
rewriter.create<LLVM::MulOp>(loc, i32Ty, leadingDim32, srcOffsetIVal);
Value loadOffset =
rewriter.create<LLVM::AddOp>(loc, i32Ty, numElemsLeadDim, srcOffsetJVal);
Value promotedSrcOpToUse;
promotedSrcOpToUse = promotedSrcOp.offset(rewriter, loc);
Value actualOffset =
rewriter.create<LLVM::AddOp>(loc, i32Ty, loadOffset, promotedSrcOpToUse);
Value loadAddress = rewriter.create<LLVM::GEPOp>(
loc,
LLVM::LLVMPointerType::get(f16Ty, srcMemrefType.getMemorySpaceAsInt()),
promotedSrcOp.alignedPtr(rewriter, loc), ArrayRef<Value>{actualOffset});
// Bitcast the base address pointer of the destination memref, So that
// values can be stored in chunks of 32-bits and semantics match with the
// intrinsic exposed by NVPTX backend.
Value loadAddressCasted = rewriter.create<LLVM::BitcastOp>(
loc,
LLVM::LLVMPointerType::get(i32Ty, srcMemrefType.getMemorySpaceAsInt()),
loadAddress);
// Get the shape of the MMAMatrix type being returned. The shape will
// choose which intrinsic this op will be lowered to.
gpu::MMAMatrixType retType =
subgroupMmaLoadMatrixOp.res().getType().cast<gpu::MMAMatrixType>();
ArrayRef<int64_t> retTypeShape = retType.getShape();
Type resType;
StringRef operandStr = retType.getOperand();
if (operandStr.equals("AOp") || operandStr.equals("BOp")) {
resType = fragArrayABTy;
} else {
if (srcMemrefType.getElementType().isF16())
resType = fragArrayCDTy;
else if (srcMemrefType.getElementType().isF32())
resType = fragArrayCDF32Ty;
else
return failure();
}
// Create nvvm.mma_load op according to the operand types.
SmallVector<Value, 2> loadOpOperands({loadAddressCasted, leadingDim32});
if (operandStr.equals("AOp")) {
if (retTypeShape[0] == 16 && retTypeShape[1] == 16) {
NVVM::WMMALoadAM16N16K16Op wmmaLoadAOp =
rewriter.create<NVVM::WMMALoadAM16N16K16Op>(loc, resType,
loadOpOperands);
rewriter.replaceOp(op, wmmaLoadAOp.getResult());
} else {
// Only M16N16K16 version implemented. Add cases as more versions are
// implemented.
return failure();
}
} else if (operandStr.equals("BOp")) {
if (retTypeShape[0] == 16 && retTypeShape[1] == 16) {
NVVM::WMMALoadBM16N16K16Op wmmaLoadBOp =
rewriter.create<NVVM::WMMALoadBM16N16K16Op>(loc, resType,
loadOpOperands);
rewriter.replaceOp(op, wmmaLoadBOp.getResult());
} else {
// Only M16N16K16 version implemented. Add cases as more versions are
// implemented.
return failure();
}
ftynseUnsubmitted
Done
ReplyInline Actions

Nit: can you put this comment inside return rewriter.notifyMatchFailure("") instead? Here and below.

ftynse: Nit: can you put this comment inside `return rewriter.notifyMatchFailure("")` instead? Here and…
} else {
if (retTypeShape[0] == 16 && retTypeShape[1] == 16) {
if (srcMemrefType.getElementType().isF16()) {
NVVM::WMMALoadCF16M16N16K16Op wmmaLoadCOp =
rewriter.create<NVVM::WMMALoadCF16M16N16K16Op>(loc, resType,
loadOpOperands);
rewriter.replaceOp(op, wmmaLoadCOp.getResult());
} else if (srcMemrefType.getElementType().isF32()) {
NVVM::WMMALoadCF32M16N16K16Op wmmaLoadCOp =
rewriter.create<NVVM::WMMALoadCF32M16N16K16Op>(loc, resType,
loadOpOperands);
rewriter.replaceOp(op, wmmaLoadCOp.getResult());
}
} else {
// Only M16N16K16 version implemented. Add cases as more versions are
// implemented.
return failure();
}
}
return success();
}
WmmaStoreOpToNVVMLowering::WmmaStoreOpToNVVMLowering(
LLVMTypeConverter &typeConverter)
: ConvertOpToLLVMPattern<gpu::SubgroupMmaStoreMatrixOp>(typeConverter),
CommonLLVMAndBuiltInMLIRTypes(&this->getTypeConverter()->getContext()) {}
LogicalResult WmmaStoreOpToNVVMLowering::matchAndRewrite(
gpu::SubgroupMmaStoreMatrixOp subgroupMmaStoreMatrixOp,
ArrayRef<Value> operands, ConversionPatternRewriter &rewriter) const {
Operation *op = subgroupMmaStoreMatrixOp.getOperation();
if (failed(areAllLLVMTypes(op, operands, rewriter)))
return failure();
int8_t indexTypeBitwidth = this->getTypeConverter()->getIndexTypeBitwidth();
// The corresponding intrinsics expects leadDimension to be a 32-bit
// integer, so all the calculations of linearizing the store address
// must also follow this restriction.
if (indexTypeBitwidth != 32)
return rewriter.notifyMatchFailure(
op, "expected indices to the memref to be 32-bit wide.");
Location loc = op->getLoc();
// Destination memref of the original op.
MemRefType dstMemrefType =
subgroupMmaStoreMatrixOp.dstMemref().getType().cast<MemRefType>();
// MemRefDescriptor to extract alignedPtr and offset.
MemRefDescriptor promotedDstOp(
gpu::SubgroupMmaStoreMatrixOpAdaptor(operands).dstMemref());
auto leadDimension = subgroupMmaStoreMatrixOp.leadDimensionAttr();
// Emit ops which compute the store offset using `dstOffsetI`,
// `dstOffsetJ`. The actualOffset is (memrefOffset + (alignedPtr +
// ((leadDimension * dstOffsetI) + dstOffsetJ)).
SmallVector<Value> indices(subgroupMmaStoreMatrixOp.indices());
Value dstOffsetIVal = indices[0];
Value dstOffsetJVal = indices[1];
Value leadingDim32 =
rewriter.create<LLVM::ConstantOp>(loc, i32Ty, leadDimension);
Value numElemsLeadDim =
rewriter.create<LLVM::MulOp>(loc, i32Ty, leadingDim32, dstOffsetIVal);
Value loadOffset =
rewriter.create<LLVM::AddOp>(loc, i32Ty, numElemsLeadDim, dstOffsetJVal);
Value promotedDstOpToUse;
promotedDstOpToUse = promotedDstOp.offset(rewriter, loc);
Value actualOffset =
rewriter.create<LLVM::AddOp>(loc, i32Ty, loadOffset, promotedDstOpToUse);
Value storeAddress = rewriter.create<LLVM::GEPOp>(
loc,
LLVM::LLVMPointerType::get(f16Ty, dstMemrefType.getMemorySpaceAsInt()),
promotedDstOp.alignedPtr(rewriter, loc), ArrayRef<Value>{actualOffset});
// Bitcast the base address pointer of the destination memref, So that
// values can be stored in chunks of 32-bits and semantics match with the
// intrinsic exposed by NVPTX backend.
Value storeAddressCasted = rewriter.create<LLVM::BitcastOp>(
loc,
LLVM::LLVMPointerType::get(i32Ty, dstMemrefType.getMemorySpaceAsInt()),
storeAddress);
SmallVector<Value, 4> storeOpOperands;
storeOpOperands.push_back(storeAddressCasted);
// Get the shape of the MMAMatrix type being stored. The shape will
// choose which intrinsic this op will be lowered to.
gpu::MMAMatrixType srcType =
subgroupMmaStoreMatrixOp.src().getType().cast<gpu::MMAMatrixType>();
ArrayRef<int64_t> srcTypeShape = srcType.getShape();
// Unpack the results from the source.
if (subgroupMmaStoreMatrixOp.src()
.getType()
.cast<gpu::MMAMatrixType>()
.getElementType() == f16Ty) {
for (unsigned i = 0, e = numHalfsInOpFrags[D]; i < e; ++i) {
Value toUse = rewriter.create<LLVM::ExtractValueOp>(
loc, f16x2Ty, operands[0], rewriter.getI32ArrayAttr(i));
storeOpOperands.push_back(toUse);
}
storeOpOperands.push_back(leadingDim32);
// Create nvvm.mma_store op.
if (srcTypeShape[0] == 16 && srcTypeShape[1] == 16) {
rewriter.create<NVVM::WMMAStoreF16M16N16K16Op>(loc, storeOpOperands);
} else {
// Only M16N16K16 version implemented. Add cases as more versions are
// implemented.
return failure();
}
rewriter.eraseOp(op);
return success();
} else if (subgroupMmaStoreMatrixOp.src()
.getType()
.cast<gpu::MMAMatrixType>()
.getElementType() == f32Ty) {
for (unsigned i = 0, e = 8; i < e; ++i) {
Value toUse = rewriter.create<LLVM::ExtractValueOp>(
loc, f32Ty, operands[0], rewriter.getI32ArrayAttr(i));
storeOpOperands.push_back(toUse);
}
storeOpOperands.push_back(leadingDim32);
// Create nvvm.mma_store op.
if (srcTypeShape[0] == 16 && srcTypeShape[1] == 16)
rewriter.create<NVVM::WMMAStoreF32M16N16K16Op>(loc, storeOpOperands);
else {
// Only M16N16K16 version implemented. Add cases as more versions are
// implemented.
return failure();
}
rewriter.eraseOp(op);
return success();
}
return failure();
}
WmmaMmaOpToNVVMLowering::WmmaMmaOpToNVVMLowering(
LLVMTypeConverter &typeConverter)
: ConvertOpToLLVMPattern<gpu::SubgroupMmaComputeOp>(typeConverter),
CommonLLVMAndBuiltInMLIRTypes(&this->getTypeConverter()->getContext()) {}
LogicalResult WmmaMmaOpToNVVMLowering::matchAndRewrite(
gpu::SubgroupMmaComputeOp subgroupMmaComputeOp, ArrayRef<Value> operands,
ConversionPatternRewriter &rewriter) const {
Operation *op = subgroupMmaComputeOp.getOperation();
if (failed(areAllLLVMTypes(op, operands, rewriter)))
return failure();
Location loc = op->getLoc();
// The wmma.mma intrinsic in llvm requires the operands as individual
// values. So individual elements from the memrefs need to be extracted and
// then passed on to the intrinsic call. Emit llvm ops to extract individual
// values form lowered memrefs.
SmallVector<Value> unpackedOps;
auto unpackOp = [&](CommonLLVMAndBuiltInMLIRTypes::OperandMap op,
Value operand, unsigned numElems, Type elemType) {
for (unsigned i = 0; i < numElems; ++i) {
Value toUse = rewriter.create<LLVM::ExtractValueOp>(
loc, elemType, operand, rewriter.getI32ArrayAttr(i));
unpackedOps.push_back(toUse);
}
};
// Get the shapes of the MMAMatrix type being used. The shapes will
// choose which intrinsic this op will be lowered to.
gpu::MMAMatrixType aType =
subgroupMmaComputeOp.opA().getType().cast<gpu::MMAMatrixType>();
ArrayRef<int64_t> aTypeShape = aType.getShape();
gpu::MMAMatrixType bType =
subgroupMmaComputeOp.opA().getType().cast<gpu::MMAMatrixType>();
ArrayRef<int64_t> bTypeShape = bType.getShape();
gpu::MMAMatrixType cType =
subgroupMmaComputeOp.opA().getType().cast<gpu::MMAMatrixType>();
ArrayRef<int64_t> cTypeShape = cType.getShape();
gpu::SubgroupMmaComputeOpAdaptor transformedOperands(operands);
if (subgroupMmaComputeOp.opC()
.getType()
.cast<gpu::MMAMatrixType>()
.getElementType() == f16Ty) {
unpackOp(A, transformedOperands.opA(), numHalfsInOpFrags[A], f16x2Ty);
unpackOp(B, transformedOperands.opB(), numHalfsInOpFrags[B], f16x2Ty);
unpackOp(C, transformedOperands.opC(), numHalfsInOpFrags[C], f16x2Ty);
if (aTypeShape[0] == 16 && aTypeShape[1] == 16 && bTypeShape[0] == 16 &&
bTypeShape[1] == 16 && cTypeShape[0] == 16 && cTypeShape[1] == 16) {
// Create nvvm.wmma.mma op.
NVVM::WMMAMmaF16F16M16N16K16Op wmmaMmaOp =
rewriter.create<NVVM::WMMAMmaF16F16M16N16K16Op>(loc, fragArrayCDTy,
unpackedOps);
rewriter.replaceOp(op, wmmaMmaOp.getResult());
return success();
} else {
// Only M16N16K16 version implemented. Add cases as more versions are
// implemented.
return failure();
}
} else if (subgroupMmaComputeOp.opC()
.getType()
.cast<gpu::MMAMatrixType>()
.getElementType() == f32Ty) {
unpackOp(A, transformedOperands.opA(), numHalfsInOpFrags[A], f16x2Ty);
unpackOp(B, transformedOperands.opB(), numHalfsInOpFrags[B], f16x2Ty);
unpackOp(C, transformedOperands.opC(), 8, f32Ty);
if (aTypeShape[0] == 16 && aTypeShape[1] == 16 && bTypeShape[0] == 16 &&
bTypeShape[1] == 16 && cTypeShape[0] == 16 && cTypeShape[1] == 16) {
// Create nvvm.wmma.mma op.
NVVM::WMMAMmaF32F32M16N16K16Op wmmaMmaOp =
rewriter.create<NVVM::WMMAMmaF32F32M16N16K16Op>(loc, fragArrayCDF32Ty,
unpackedOps);
rewriter.replaceOp(op, wmmaMmaOp.getResult());
return success();
} else {
// Only M16N16K16 version implemented. Add cases as more versions are
// implemented.
return failure();
}
}
return failure();
}
} // namespace mlir