lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp - llvm-project/mlir - Git at Google

 //===------ 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);
 }
	//===------ 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);
	}