Page MenuHomePhabricator
Differential D95331 Diff 343284 mlir/lib/Conversion/GPUToNVVM/WmmaMmaOptoNvvmLowering.h

Changeset View

Standalone View

View Options

mlir/lib/Conversion/GPUToNVVM/WmmaMmaOptoNvvmLowering.h

  • This file was added.
//===--- WmmaMmaOpToNVVMLowering.h - MmaOp to NVVM Op lowering -*- C++ -*-===//
//
// 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 patterns to lower the GPU subgroup MMA_computeOp to the
// NVVM Dialect.
//
//===----------------------------------------------------------------------===//
#ifndef MLIR_LIB_CONVERSION_GPUTONVVM_WMMASTOREOPTONVVMLOWERING_H
Lint: Pre-merge checks
Inline Actions

clang-tidy: warning: header guard does not follow preferred style [llvm-header-guard]
not useful

Lint: Pre-merge checks: clang-tidy: warning: header guard does not follow preferred style [llvm-header-guard] [[https…
#define MLIR_LIB_CONVERSION_GPUTONVVM_WMMASTOREOPTONVVMLOWERING_H
#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 {
/// This class implemtents the conversion of GPU MMA computeOp to wmma.mma op
/// in the NVVM dialect.
struct WmmaMmaOpToNVVMLowering
: public ConvertOpToLLVMPattern<gpu::SubgroupMmaComputeOp> {
public:
MLIRContext *context = &this->getTypeConverter()->getContext();
explicit WmmaMmaOpToNVVMLowering(LLVMTypeConverter &typeConverter)
: ConvertOpToLLVMPattern<gpu::SubgroupMmaComputeOp>(typeConverter),
llvmMLIRTypes(context) {}
static LogicalResult areAllLLVMTypes(Operation *op, ValueRange operands,
ConversionPatternRewriter &rewriter) {
if (!llvm::all_of(operands, [](Value value) {
return LLVM::isCompatibleType(value.getType());
}))
return rewriter.notifyMatchFailure(
op, "Cannot convert if operands aren't of LLVM type.");
return success();
}
ftynseUnsubmitted
Done
ReplyInline Actions

A similar function is present in another file this patch is adding. Consider refactoring to only have one definition.

ftynse: A similar function is present in another file this patch is adding. Consider refactoring to…
LogicalResult
matchAndRewrite(gpu::SubgroupMmaComputeOp subgroupMmaComputeOp,
ArrayRef<Value> operands,
ConversionPatternRewriter &rewriter) const override {
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, 24> unpackedOps;
ftynseUnsubmitted
Done
ReplyInline Actions

SmallVector now has a default number of stack elements. Drop the manually specified number unless there is a specific reason to choose the value. Here, I don't see why 24 is special.

ftynse: SmallVector now has a default number of stack elements. Drop the manually specified number…
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();
if (subgroupMmaComputeOp.opC()
.getType()
.cast<gpu::MMAMatrixType>()
.getElementType() == llvmMLIRTypes.f16Ty) {
unpackOp(llvmMLIRTypes.A, operands[0],
ftynseUnsubmitted
Done
ReplyInline Actions

Please use OpAdaptors here and below to avoid operands[magic-constant]

ftynse: Please use OpAdaptors here and below to avoid `operands[magic-constant]`
llvmMLIRTypes.numHalfsInOpFrags[llvmMLIRTypes.A],
llvmMLIRTypes.f16x2Ty);
unpackOp(llvmMLIRTypes.B, operands[1],
llvmMLIRTypes.numHalfsInOpFrags[llvmMLIRTypes.B],
llvmMLIRTypes.f16x2Ty);
unpackOp(llvmMLIRTypes.C, operands[2],
llvmMLIRTypes.numHalfsInOpFrags[llvmMLIRTypes.C],
llvmMLIRTypes.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, llvmMLIRTypes.fragArrayCDTy, unpackedOps);
rewriter.replaceOp(op, wmmaMmaOp.getResult());
return success();
} else {
// Only M16N16K16 version implemented. Add cases as more versions are
// implemented.
return failure();
}
}
if (subgroupMmaComputeOp.opC()
.getType()
.cast<gpu::MMAMatrixType>()
.getElementType() == llvmMLIRTypes.f32Ty) {
unpackOp(llvmMLIRTypes.A, operands[0],
llvmMLIRTypes.numHalfsInOpFrags[llvmMLIRTypes.A],
llvmMLIRTypes.f16x2Ty);
unpackOp(llvmMLIRTypes.B, operands[1],
llvmMLIRTypes.numHalfsInOpFrags[llvmMLIRTypes.B],
llvmMLIRTypes.f16x2Ty);
unpackOp(llvmMLIRTypes.C, operands[2], 8, llvmMLIRTypes.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, llvmMLIRTypes.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();
}
private:
/// Definitions of all the LLVM types which are used for lowering
/// this GPU subgroupMmaComputeOp.
CommonLLVMAndBuiltInMLIRTypes llvmMLIRTypes;
};
} // namespace mlir
#endif