llvm/mlir/lib/Dialect/AMX/Transforms/LegalizeForLLVMExport.cpp

//===- LegalizeForLLVMExport.cpp - Prepare AMX for LLVM translation ----===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/AMX/Transforms.h"

#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Dialect/AMX/AMXDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/PatternMatch.h"

usingnamespacemlir;
usingnamespacemlir::amx;

namespace {

/// Maps the 2-dim vector shape to the two 16-bit tile sizes. The first
/// dimension directly translates into the number of rows of the tiles.
/// The second dimensions needs to be scaled by the number of bytes.
std::pair<Value, Value> getTileSizes(ConversionPatternRewriter &rewriter,
                                     const LLVMTypeConverter &typeConverter,
                                     VectorType vType, Location loc) { … }

/// Verifies if the stride matches proper tile access.
LogicalResult verifyStride(MemRefType mType) { … }

/// Maps the 2-dim memref shape to the 64-bit stride. Note that the buffer
/// shape may "envelop" the actual tile shape, and may be dynamically sized.
Value getStride(ConversionPatternRewriter &rewriter,
                const LLVMTypeConverter &typeConverter, MemRefType mType,
                Value base, Location loc) { … }

struct TileZeroConversion : public ConvertOpToLLVMPattern<TileZeroOp> { … };

struct TileLoadConversion : public ConvertOpToLLVMPattern<TileLoadOp> { … };

struct TileStoreConversion : public ConvertOpToLLVMPattern<TileStoreOp> { … };

struct TileMulFConversion : public ConvertOpToLLVMPattern<TileMulFOp> { … };

struct TileMulIConversion : public ConvertOpToLLVMPattern<TileMulIOp> { … };

} // namespace

void mlir::populateAMXLegalizeForLLVMExportPatterns(
    const LLVMTypeConverter &converter, RewritePatternSet &patterns) { … }

void mlir::configureAMXLegalizeForExportTarget(LLVMConversionTarget &target) { … }