71 lines · cpp
1//===- LegalizeForLLVMExport.cpp - Prepare AMX for LLVM translation ----===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/AMX/Transforms.h"10 11#include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"12#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"13#include "mlir/Conversion/LLVMCommon/Pattern.h"14#include "mlir/Dialect/AMX/AMXDialect.h"15#include "mlir/IR/PatternMatch.h"16 17using namespace mlir;18using namespace mlir::amx;19 20namespace {21 22/// Generic one-to-one conversion of simply mappable operations into calls23/// to their respective LLVM intrinsics.24struct AMXIntrinsicOpConversion25 : public ConvertOpInterfaceToLLVMPattern<amx::AMXIntrinsicOp> {26 using ConvertOpInterfaceToLLVMPattern::ConvertOpInterfaceToLLVMPattern;27 28 LogicalResult29 matchAndRewrite(amx::AMXIntrinsicOp op, ArrayRef<Value> operands,30 ConversionPatternRewriter &rewriter) const override {31 const LLVMTypeConverter &typeConverter = *getTypeConverter();32 return LLVM::detail::intrinsicRewrite(33 op, rewriter.getStringAttr(op.getIntrinsicName()),34 op.getIntrinsicOperands(operands, typeConverter, rewriter),35 typeConverter, rewriter);36 }37};38 39} // namespace40 41void mlir::populateAMXLegalizeForLLVMExportPatterns(42 LLVMTypeConverter &converter, RewritePatternSet &patterns) {43 patterns.add<AMXIntrinsicOpConversion>(converter);44 converter.addConversion([&](amx::TileType type) {45 return LLVM::LLVMX86AMXType::get(&converter.getContext());46 });47}48 49void mlir::configureAMXLegalizeForExportTarget(LLVMConversionTarget &target) {50 target.addIllegalDialect<AMXDialect>();51}52 53namespace {54/// Implement the interface to convert AMX to LLVM.55struct AMXToLLVMDialectInterface : public ConvertToLLVMPatternInterface {56 using ConvertToLLVMPatternInterface::ConvertToLLVMPatternInterface;57 58 void populateConvertToLLVMConversionPatterns(59 ConversionTarget &target, LLVMTypeConverter &typeConverter,60 RewritePatternSet &patterns) const final {61 populateAMXLegalizeForLLVMExportPatterns(typeConverter, patterns);62 }63};64} // namespace65 66void mlir::registerConvertAMXToLLVMInterface(DialectRegistry ®istry) {67 registry.addExtension(+[](MLIRContext *ctx, amx::AMXDialect *dialect) {68 dialect->addInterfaces<AMXToLLVMDialectInterface>();69 });70}71