WmmaOpsToSPIRV.cpp

Go to the documentation of this file.

//===------ WmmaOpsToSPIRV.cpp - WMMA LD/ST/Compute to SPIRV 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
// SPIRV Cooperative Matrix ops.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/ValueRange.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/StringSwitch.h"
 
#include <cassert>
#include <limits>
 
namespace mlir {
//===----------------------------------------------------------------------===//
// Patterns and helpers.
//===----------------------------------------------------------------------===//
 
static spirv::CooperativeMatrixType
convertMMAMatrixType(gpu::MMAMatrixType type) {
  ArrayRef<int64_t> shape = type.getShape();
  auto use =
      llvm::StringSwitch<spirv::CooperativeMatrixUseKHR>(type.getOperand())
          .Case("AOp", spirv::CooperativeMatrixUseKHR::MatrixA)
          .Case("BOp", spirv::CooperativeMatrixUseKHR::MatrixB)
          .Default(spirv::CooperativeMatrixUseKHR::MatrixAcc);
  return spirv::CooperativeMatrixType::get(
      type.getElementType(), shape[0], shape[1], spirv::Scope::Subgroup, use);
}
 
// Convert a memref of `gpu.mma_matrix` into a SPIR-V pointer to an array
// of spirv::CooperativeMatrix.
static std::optional<Type> convertMemrefOfMMAMatrixType(MemRefType type) {
  auto matrixType = dyn_cast<gpu::MMAMatrixType>(type.getElementType());
  if (!matrixType)
    return std::nullopt;
  // SPV Cooperative matrix types, and composites containing them, may only be
  // allocated in Function or Private storage classes. For now, support only
  // function-local arrays.
  auto storageClass =
      dyn_cast_or_null<spirv::StorageClassAttr>(type.getMemorySpace());
  if (!storageClass ||
      storageClass.getValue() != spirv::StorageClass::Function ||
      !type.hasStaticShape() || !type.getLayout().isIdentity() ||
      type.getNumElements() <= 0 ||
      type.getNumElements() > std::numeric_limits<unsigned>::max())
    return Type();
  auto arrayType =
      spirv::ArrayType::get(convertMMAMatrixType(matrixType),
                            static_cast<unsigned>(type.getNumElements()));
  return spirv::PointerType::get(arrayType, spirv::StorageClass::Function);
}
 
/// Creates a SPIR-V op to replace the given GPU subgroup mma elementwise op
/// when the elementwise op directly supports with cooperative matrix type.
/// Returns false if cannot.
///
/// See SPV_KHR_cooperative_matrix for supported elementwise ops.
static bool createElementwiseOp(ConversionPatternRewriter &builder,
                                gpu::SubgroupMmaElementwiseOp op, Type coopType,
                                ValueRange operands) {
  assert((isa<spirv::CooperativeMatrixType>(coopType)));
 
  switch (op.getOpType()) {
  case gpu::MMAElementwiseOp::ADDF:
    builder.replaceOpWithNewOp<spirv::FAddOp>(op, coopType, operands);
    return true;
  case gpu::MMAElementwiseOp::ADDI:
    builder.replaceOpWithNewOp<spirv::IAddOp>(op, coopType, operands);
    return true;
  case gpu::MMAElementwiseOp::SUBF:
    builder.replaceOpWithNewOp<spirv::FSubOp>(op, coopType, operands);
    return true;
  case gpu::MMAElementwiseOp::SUBI:
    builder.replaceOpWithNewOp<spirv::ISubOp>(op, coopType, operands);
    return true;
  case gpu::MMAElementwiseOp::MULF:
    builder.replaceOpWithNewOp<spirv::FMulOp>(op, coopType, operands);
    return true;
  case gpu::MMAElementwiseOp::DIVF:
    builder.replaceOpWithNewOp<spirv::FDivOp>(op, coopType, operands);
    return true;
  case gpu::MMAElementwiseOp::DIVS:
    builder.replaceOpWithNewOp<spirv::SDivOp>(op, coopType, operands);
    return true;
  case gpu::MMAElementwiseOp::DIVU:
    builder.replaceOpWithNewOp<spirv::UDivOp>(op, coopType, operands);
    return true;
  case gpu::MMAElementwiseOp::NEGATEF:
    builder.replaceOpWithNewOp<spirv::FNegateOp>(op, coopType, operands);
    return true;
  case gpu::MMAElementwiseOp::NEGATES:
    builder.replaceOpWithNewOp<spirv::SNegateOp>(op, coopType, operands);
    return true;
  case gpu::MMAElementwiseOp::EXTF:
  case gpu::MMAElementwiseOp::TRUNCF:
    builder.replaceOpWithNewOp<spirv::FConvertOp>(op, coopType, operands);
    return true;
  default:
    break;
  }
  return false;
}
 
bool allOperandsHaveSameCoopMatrixType(ValueRange operands) {
  assert(!operands.empty());
  if (!llvm::all_equal(
          llvm::map_range(operands, [](Value v) { return v.getType(); })))
    return false;
 
  return isa<spirv::CooperativeMatrixType>(operands.front().getType());
}
 
static bool hasSignedIntegerElementType(spirv::CooperativeMatrixType type) {
  auto elementType = dyn_cast<IntegerType>(type.getElementType());
  return elementType && elementType.isSigned();
}
 
static spirv::CooperativeMatrixOperandsKHR
getSignedCoopMatrixOperands(spirv::CooperativeMatrixType aType,
                            spirv::CooperativeMatrixType bType,
                            spirv::CooperativeMatrixType cType,
                            spirv::CooperativeMatrixType resultType) {
  using Operands = spirv::CooperativeMatrixOperandsKHR;
 
  Operands operands = Operands::None;
  if (hasSignedIntegerElementType(aType))
    operands |= Operands::ASigned;
  if (hasSignedIntegerElementType(bType))
    operands |= Operands::BSigned;
  if (hasSignedIntegerElementType(cType))
    operands |= Operands::CSigned;
  if (hasSignedIntegerElementType(resultType))
    operands |= Operands::ResultSigned;
  return operands;
}
 
namespace {
/// Converts GPU MMA ConstantMatrixOp to constant SPIR-V KHR/NV cooperative
/// matrix ops.
struct WmmaConstantOpToSPIRVLowering final
    : OpConversionPattern<gpu::SubgroupMmaConstantMatrixOp> {
  using Base::Base;
 
  LogicalResult
  matchAndRewrite(gpu::SubgroupMmaConstantMatrixOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value cst = llvm::getSingleElement(adaptor.getOperands());
    auto coopType = getTypeConverter()->convertType(op.getType());
    if (!coopType)
      return rewriter.notifyMatchFailure(op, "type conversion failed");
 
    rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(op, coopType, cst);
    return success();
  }
};
 
/// Converts GPU MMA ExtractOp to CompositeExtract SPIR-V KHR/NV cooperative
/// matrix ops.
struct WmmaExtractOpToSPIRVLowering final
    : OpConversionPattern<gpu::SubgroupMmaExtractThreadLocalOp> {
  using Base::Base;
 
  LogicalResult
  matchAndRewrite(gpu::SubgroupMmaExtractThreadLocalOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value matrix = adaptor.getMatrix();
    auto coopType =
        getTypeConverter()->convertType<spirv::CooperativeMatrixType>(
            matrix.getType());
    if (!coopType)
      return rewriter.notifyMatchFailure(op, "type conversion failed");
 
    SmallVector<int32_t> intValues;
    for (Value val : op.getIndices()) {
      if (auto constOp = val.getDefiningOp<arith::ConstantIndexOp>()) {
        intValues.push_back(static_cast<int32_t>(constOp.value()));
      } else {
        return rewriter.notifyMatchFailure(op, "indices must be constants");
      }
    }
 
    Type elementType = coopType.getElementType();
    rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
        op, elementType, matrix, rewriter.getI32ArrayAttr(intValues));
    return success();
  }
};
 
/// Converts GPU MMA InsertOp to CompositeInsert SPIR-V KHR/NV cooperative
/// matrix ops.
struct WmmaInsertOpToSPIRVLowering final
    : OpConversionPattern<gpu::SubgroupMmaInsertThreadLocalOp> {
  using Base::Base;
 
  LogicalResult
  matchAndRewrite(gpu::SubgroupMmaInsertThreadLocalOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value value = adaptor.getValue();
    Value matrix = adaptor.getMatrix();
    auto coopType = getTypeConverter()->convertType(matrix.getType());
    if (!coopType)
      return rewriter.notifyMatchFailure(op, "type conversion failed");
 
    SmallVector<int32_t> intValues;
    for (Value val : op.getIndices()) {
      if (auto constOp = val.getDefiningOp<arith::ConstantIndexOp>()) {
        intValues.push_back(static_cast<int32_t>(constOp.value()));
      } else {
        return rewriter.notifyMatchFailure(op, "indices must be constants");
      }
    }
 
    rewriter.replaceOpWithNewOp<spirv::CompositeInsertOp>(
        op, coopType, value, matrix, rewriter.getI32ArrayAttr(intValues));
    return success();
  }
};
 
/// Converts elementwise ops to SPIR-V cooperative matrix elementwise ops for
/// the default case.
struct WmmaElementwiseOpToSPIRVDefaultLowering final
    : OpConversionPattern<gpu::SubgroupMmaElementwiseOp> {
  using Base::Base;
 
  LogicalResult
  matchAndRewrite(gpu::SubgroupMmaElementwiseOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // All operands should be of cooperative matrix types.
    if (!allOperandsHaveSameCoopMatrixType(adaptor.getOperands())) {
      return rewriter.notifyMatchFailure(op,
                                         "not all operands are coop matrices");
    }
 
    auto coopType = getTypeConverter()->convertType(op.getType());
    if (!coopType)
      return rewriter.notifyMatchFailure(op, "type conversion failed");
 
    return success(
        createElementwiseOp(rewriter, op, coopType, adaptor.getOperands()));
  }
};
 
/// Converts elementwise ops to SPIR-V cooperative matrix elementwise ops for
/// matrix times scalar case.
struct WmmaElementwiseOpToSPIRVScalarMulLowering final
    : OpConversionPattern<gpu::SubgroupMmaElementwiseOp> {
  using Base::Base;
 
  LogicalResult
  matchAndRewrite(gpu::SubgroupMmaElementwiseOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (adaptor.getOperands().size() != 2)
      return failure();
 
    // All operands should be of cooperative matrix types.
    if (!allOperandsHaveSameCoopMatrixType(adaptor.getOperands())) {
      return rewriter.notifyMatchFailure(op,
                                         "not all operands are coop matrices");
    }
 
    if (op.getOpType() != gpu::MMAElementwiseOp::MULF)
      return failure();
 
    // Use the original operands to check whether one of the operands is a splat
    // scalar value.
    Value lhs = op.getOperands().front();
    Value rhs = op.getOperands().back();
    Value splat = nullptr;
    Value matrix = nullptr;
    if (lhs.getDefiningOp<gpu::SubgroupMmaConstantMatrixOp>()) {
      splat = adaptor.getOperands().front();
      matrix = adaptor.getOperands().back();
    } else if (rhs.getDefiningOp<gpu::SubgroupMmaConstantMatrixOp>()) {
      matrix = adaptor.getOperands().front();
      splat = adaptor.getOperands().back();
    }
    if (!splat || !matrix)
      return rewriter.notifyMatchFailure(op, "no splat operand");
 
    // Constant MMA matrix ops are converted to `spirv.CompositeConstruct` ops.
    Value scalar;
    auto cc = splat.getDefiningOp<spirv::CompositeConstructOp>();
    if (!cc) {
      return rewriter.notifyMatchFailure(op,
                                         "splat is not a composite construct");
    }
 
    scalar = llvm::getSingleElement(cc.getConstituents());
 
    auto coopType = getTypeConverter()->convertType(op.getType());
    if (!coopType)
      return rewriter.notifyMatchFailure(op, "type conversion failed");
    rewriter.replaceOpWithNewOp<spirv::MatrixTimesScalarOp>(
        op, coopType, ValueRange{matrix, scalar});
    return success();
  }
};
} // namespace
 
//===----------------------------------------------------------------------===//
// SPV_KHR_cooperative_matrix
//===----------------------------------------------------------------------===//
 
namespace khr {
namespace {
 
/// Converts the GPU MMA loadOp to KHRCooperativeMatrixLoad op in the SPIRV
/// dialect.
struct WmmaLoadOpToSPIRVLowering final
    : OpConversionPattern<gpu::SubgroupMmaLoadMatrixOp> {
  using Base::Base;
 
  LogicalResult
  matchAndRewrite(gpu::SubgroupMmaLoadMatrixOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    const auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
    Location loc = op->getLoc();
 
    auto retType = cast<gpu::MMAMatrixType>(op.getRes().getType());
    MemRefType memrefType = op.getSrcMemref().getType();
    Value bufferPtr =
        spirv::getElementPtr(typeConverter, memrefType, adaptor.getSrcMemref(),
                             adaptor.getIndices(), loc, rewriter);
 
    auto coopType =
        typeConverter.convertType<spirv::CooperativeMatrixType>(retType);
    if (!coopType)
      return rewriter.notifyMatchFailure(op, "type conversion failed");
 
    int64_t stride = op.getLeadDimension().getSExtValue();
    IntegerType i32Type = rewriter.getI32Type();
    auto strideValue = spirv::ConstantOp::create(
        rewriter, loc, i32Type, IntegerAttr::get(i32Type, stride));
 
    bool isColMajor = op.getTranspose().value_or(false);
    auto layout = isColMajor ? spirv::CooperativeMatrixLayoutKHR::ColumnMajor
                             : spirv::CooperativeMatrixLayoutKHR::RowMajor;
 
    rewriter.replaceOpWithNewOp<spirv::KHRCooperativeMatrixLoadOp>(
        op, coopType, bufferPtr, strideValue, layout);
    return success();
  }
};
 
/// Converts the GPU MMA StoreOp to KHRCooperativeMatrixStore op in the SPIRV
/// dialect.
struct WmmaStoreOpToSPIRVLowering final
    : OpConversionPattern<gpu::SubgroupMmaStoreMatrixOp> {
  using Base::Base;
 
  LogicalResult
  matchAndRewrite(gpu::SubgroupMmaStoreMatrixOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    const auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
    Location loc = op->getLoc();
 
    auto memrefType = cast<MemRefType>(op.getDstMemref().getType());
    Value bufferPtr =
        spirv::getElementPtr(typeConverter, memrefType, adaptor.getDstMemref(),
                             adaptor.getIndices(), loc, rewriter);
 
    int64_t stride = op.getLeadDimension().getSExtValue();
    IntegerType i32Type = rewriter.getI32Type();
    auto strideValue = spirv::ConstantOp::create(
        rewriter, loc, i32Type, IntegerAttr::get(i32Type, stride));
 
    bool isColMajor = op.getTranspose().value_or(false);
    auto layout = isColMajor ? spirv::CooperativeMatrixLayoutKHR::ColumnMajor
                             : spirv::CooperativeMatrixLayoutKHR::RowMajor;
 
    rewriter.replaceOpWithNewOp<spirv::KHRCooperativeMatrixStoreOp>(
        op, bufferPtr, adaptor.getSrc(), strideValue, layout);
    return success();
  }
};
 
/// Converts GPU MMA Compute to KHRCooperativeMatrixMulAdd op in the SPIRV
/// dialect.
struct WmmaMmaOpToSPIRVLowering final
    : OpConversionPattern<gpu::SubgroupMmaComputeOp> {
  using Base::Base;
 
  LogicalResult
  matchAndRewrite(gpu::SubgroupMmaComputeOp subgroupMmaComputeOp,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto aType =
        dyn_cast<spirv::CooperativeMatrixType>(adaptor.getOpA().getType());
    auto bType =
        dyn_cast<spirv::CooperativeMatrixType>(adaptor.getOpB().getType());
    auto cType =
        dyn_cast<spirv::CooperativeMatrixType>(adaptor.getOpC().getType());
    auto resultType =
        getTypeConverter()->convertType<spirv::CooperativeMatrixType>(
            subgroupMmaComputeOp.getResult().getType());
    if (!aType || !bType || !cType || !resultType)
      return rewriter.notifyMatchFailure(subgroupMmaComputeOp,
                                         "type conversion failed");
 
    using Operands = spirv::CooperativeMatrixOperandsKHR;
    Operands operands =
        getSignedCoopMatrixOperands(aType, bType, cType, resultType);
    spirv::CooperativeMatrixOperandsKHRAttr operandsAttr;
    if (operands != Operands::None)
      operandsAttr = spirv::CooperativeMatrixOperandsKHRAttr::get(
          rewriter.getContext(), operands);
 
    rewriter.replaceOpWithNewOp<spirv::KHRCooperativeMatrixMulAddOp>(
        subgroupMmaComputeOp, adaptor.getOpA(), adaptor.getOpB(),
        adaptor.getOpC(), operandsAttr);
    return success();
  }
};
 
} // namespace
} // namespace khr
} // namespace mlir
 
void mlir::populateGpuWMMAToSPIRVCoopMatrixKHRConversionPatterns(
    const SPIRVTypeConverter &converter, RewritePatternSet &patterns) {
  using namespace mlir;
  MLIRContext *context = patterns.getContext();
  patterns.add<khr::WmmaLoadOpToSPIRVLowering, khr::WmmaMmaOpToSPIRVLowering,
               khr::WmmaStoreOpToSPIRVLowering, WmmaConstantOpToSPIRVLowering,
               WmmaExtractOpToSPIRVLowering, WmmaInsertOpToSPIRVLowering,
               WmmaElementwiseOpToSPIRVDefaultLowering>(converter, context);
  // Give the following patterns higher benefit to prevail over the default one.
  patterns.add<WmmaElementwiseOpToSPIRVScalarMulLowering>(converter, context,
                                                          /*benefit=*/2);
}
 
void mlir::populateMMAToSPIRVCoopMatrixTypeConversion(
    mlir::SPIRVTypeConverter &typeConverter) {
  typeConverter.addConversion(convertMMAMatrixType);
  typeConverter.addConversion(convertMemrefOfMMAMatrixType);
}