//===- LinalgToLLVM.cpp - conversion from Linalg to LLVM dialect ----------===// // // Part of the MLIR 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/Conversion/LinalgToLLVM/LinalgToLLVM.h" #include "mlir/Conversion/AffineToStandard/AffineToStandard.h" #include "mlir/Conversion/LoopToStandard/ConvertLoopToStandard.h" #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h" #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h" #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/Linalg/IR/LinalgOps.h" #include "mlir/Dialect/Linalg/IR/LinalgTypes.h" #include "mlir/Dialect/Linalg/Passes.h" #include "mlir/Dialect/Linalg/Utils/Intrinsics.h" #include "mlir/EDSC/Builders.h" #include "mlir/EDSC/Intrinsics.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/Builders.h" #include "mlir/IR/MLIRContext.h" #include "mlir/IR/Module.h" #include "mlir/IR/Operation.h" #include "mlir/IR/PatternMatch.h" #include "mlir/IR/StandardTypes.h" #include "mlir/IR/Types.h" #include "mlir/Pass/Pass.h" #include "mlir/Pass/PassManager.h" #include "mlir/Support/LogicalResult.h" #include "mlir/Transforms/DialectConversion.h" #include "mlir/Transforms/Passes.h" #include "llvm/ADT/SetVector.h" #include "llvm/IR/DerivedTypes.h" #include "llvm/IR/Module.h" #include "llvm/IR/Type.h" #include "llvm/Support/Allocator.h" #include "llvm/Support/ErrorHandling.h" using namespace mlir; using namespace mlir::edsc; using namespace mlir::edsc::intrinsics; using namespace mlir::LLVM; using namespace mlir::linalg; using namespace mlir::linalg::intrinsics; using add = ValueBuilder; using addi = ValueBuilder; using bitcast = ValueBuilder; using cmpi = ValueBuilder; using constant = ValueBuilder; using extractvalue = ValueBuilder; using gep = ValueBuilder; using insertvalue = ValueBuilder; using llvm_call = OperationBuilder; using llvm_icmp = ValueBuilder; using llvm_load = ValueBuilder; using llvm_store = OperationBuilder; using llvm_select = ValueBuilder; using mul = ValueBuilder; using ptrtoint = ValueBuilder; using sub = ValueBuilder; using llvm_undef = ValueBuilder; using urem = ValueBuilder; using llvm_alloca = ValueBuilder; using llvm_return = OperationBuilder; namespace { template static LLVMType getPtrToElementType(T containerType, LLVMTypeConverter &lowering) { return lowering.convertType(containerType.getElementType()) .template cast() .getPointerTo(); } // Convert the given type to the LLVM IR Dialect type. The following // conversions are supported: // - an Index type is converted into an LLVM integer type with pointer // bitwidth (analogous to intptr_t in C); // - an Integer type is converted into an LLVM integer type of the same width; // - an F32 type is converted into an LLVM float type // - a Buffer, Range or View is converted into an LLVM structure type // containing the respective dynamic values. static Type convertLinalgType(Type t, LLVMTypeConverter &lowering) { auto *context = t.getContext(); auto int64Ty = lowering.convertType(IntegerType::get(64, context)) .cast(); // Range descriptor contains the range bounds and the step as 64-bit integers. // // struct { // int64_t min; // int64_t max; // int64_t step; // }; if (t.isa()) return LLVMType::getStructTy(int64Ty, int64Ty, int64Ty); return Type(); } /// EDSC-compatible wrapper for MemRefDescriptor. class BaseViewConversionHelper { public: BaseViewConversionHelper(Type type) : d(MemRefDescriptor::undef(rewriter(), loc(), type)) {} BaseViewConversionHelper(Value v) : d(v) {} /// Wrappers around MemRefDescriptor that use EDSC builder and location. Value allocatedPtr() { return d.allocatedPtr(rewriter(), loc()); } void setAllocatedPtr(Value v) { d.setAllocatedPtr(rewriter(), loc(), v); } Value alignedPtr() { return d.alignedPtr(rewriter(), loc()); } void setAlignedPtr(Value v) { d.setAlignedPtr(rewriter(), loc(), v); } Value offset() { return d.offset(rewriter(), loc()); } void setOffset(Value v) { d.setOffset(rewriter(), loc(), v); } Value size(unsigned i) { return d.size(rewriter(), loc(), i); } void setSize(unsigned i, Value v) { d.setSize(rewriter(), loc(), i, v); } void setConstantSize(unsigned i, int64_t v) { d.setConstantSize(rewriter(), loc(), i, v); } Value stride(unsigned i) { return d.stride(rewriter(), loc(), i); } void setStride(unsigned i, Value v) { d.setStride(rewriter(), loc(), i, v); } void setConstantStride(unsigned i, int64_t v) { d.setConstantStride(rewriter(), loc(), i, v); } operator Value() { return d; } private: OpBuilder &rewriter() { return ScopedContext::getBuilder(); } Location loc() { return ScopedContext::getLocation(); } MemRefDescriptor d; }; // RangeOp creates a new range descriptor. class RangeOpConversion : public LLVMOpLowering { public: explicit RangeOpConversion(MLIRContext *context, LLVMTypeConverter &lowering_) : LLVMOpLowering(RangeOp::getOperationName(), context, lowering_) {} PatternMatchResult matchAndRewrite(Operation *op, ArrayRef operands, ConversionPatternRewriter &rewriter) const override { auto rangeOp = cast(op); auto rangeDescriptorTy = convertLinalgType(rangeOp.getResult().getType(), lowering); edsc::ScopedContext context(rewriter, op->getLoc()); // Fill in an aggregate value of the descriptor. RangeOpOperandAdaptor adaptor(operands); Value desc = llvm_undef(rangeDescriptorTy); desc = insertvalue(desc, adaptor.min(), rewriter.getI64ArrayAttr(0)); desc = insertvalue(desc, adaptor.max(), rewriter.getI64ArrayAttr(1)); desc = insertvalue(desc, adaptor.step(), rewriter.getI64ArrayAttr(2)); rewriter.replaceOp(op, desc); return matchSuccess(); } }; // ReshapeOp creates a new view descriptor of the proper rank. // For now, the only conversion supported is for target MemRef with static sizes // and strides. class ReshapeOpConversion : public LLVMOpLowering { public: explicit ReshapeOpConversion(MLIRContext *context, LLVMTypeConverter &lowering_) : LLVMOpLowering(ReshapeOp::getOperationName(), context, lowering_) {} PatternMatchResult matchAndRewrite(Operation *op, ArrayRef operands, ConversionPatternRewriter &rewriter) const override { auto reshapeOp = cast(op); MemRefType dstType = reshapeOp.getResult().getType().cast(); if (!dstType.hasStaticShape()) return matchFailure(); int64_t offset; SmallVector strides; auto res = getStridesAndOffset(dstType, strides, offset); if (failed(res) || llvm::any_of(strides, [](int64_t val) { return ShapedType::isDynamicStrideOrOffset(val); })) return matchFailure(); edsc::ScopedContext context(rewriter, op->getLoc()); ReshapeOpOperandAdaptor adaptor(operands); BaseViewConversionHelper baseDesc(adaptor.view()); BaseViewConversionHelper desc(lowering.convertType(dstType)); desc.setAllocatedPtr(baseDesc.allocatedPtr()); desc.setAlignedPtr(baseDesc.alignedPtr()); desc.setOffset(baseDesc.offset()); for (auto en : llvm::enumerate(dstType.getShape())) desc.setConstantSize(en.index(), en.value()); for (auto en : llvm::enumerate(strides)) desc.setConstantStride(en.index(), en.value()); rewriter.replaceOp(op, {desc}); return matchSuccess(); } }; /// Conversion pattern that transforms a linalg.slice op into: /// 1. A function entry `alloca` operation to allocate a ViewDescriptor. /// 2. A load of the ViewDescriptor from the pointer allocated in 1. /// 3. Updates to the ViewDescriptor to introduce the data ptr, offset, size /// and stride corresponding to the region of memory within the bounds of /// the parent view. /// 4. A store of the resulting ViewDescriptor to the alloca'ed pointer. /// The linalg.slice op is replaced by the alloca'ed pointer. class SliceOpConversion : public LLVMOpLowering { public: explicit SliceOpConversion(MLIRContext *context, LLVMTypeConverter &lowering_) : LLVMOpLowering(SliceOp::getOperationName(), context, lowering_) {} PatternMatchResult matchAndRewrite(Operation *op, ArrayRef operands, ConversionPatternRewriter &rewriter) const override { edsc::ScopedContext context(rewriter, op->getLoc()); SliceOpOperandAdaptor adaptor(operands); BaseViewConversionHelper baseDesc(adaptor.view()); auto sliceOp = cast(op); auto memRefType = sliceOp.getBaseViewType(); auto int64Ty = lowering.convertType(rewriter.getIntegerType(64)) .cast(); BaseViewConversionHelper desc( lowering.convertType(sliceOp.getShapedType())); // TODO(ntv): extract sizes and emit asserts. SmallVector strides(memRefType.getRank()); for (int i = 0, e = memRefType.getRank(); i < e; ++i) strides[i] = baseDesc.stride(i); auto pos = [&rewriter](ArrayRef values) { return rewriter.getI64ArrayAttr(values); }; // Compute base offset. Value baseOffset = baseDesc.offset(); for (int i = 0, e = memRefType.getRank(); i < e; ++i) { Value indexing = adaptor.indexings()[i]; Value min = indexing; if (sliceOp.indexing(i).getType().isa()) min = extractvalue(int64Ty, indexing, pos(0)); baseOffset = add(baseOffset, mul(min, strides[i])); } // Insert the base and aligned pointers. desc.setAllocatedPtr(baseDesc.allocatedPtr()); desc.setAlignedPtr(baseDesc.alignedPtr()); // Insert base offset. desc.setOffset(baseOffset); // Corner case, no sizes or strides: early return the descriptor. if (sliceOp.getShapedType().getRank() == 0) return rewriter.replaceOp(op, {desc}), matchSuccess(); Value zero = constant(int64Ty, rewriter.getIntegerAttr(rewriter.getIndexType(), 0)); // Compute and insert view sizes (max - min along the range) and strides. // Skip the non-range operands as they will be projected away from the view. int numNewDims = 0; for (auto en : llvm::enumerate(sliceOp.indexings())) { Value indexing = en.value(); if (indexing.getType().isa()) { int rank = en.index(); Value rangeDescriptor = adaptor.indexings()[rank]; Value min = extractvalue(int64Ty, rangeDescriptor, pos(0)); Value max = extractvalue(int64Ty, rangeDescriptor, pos(1)); Value step = extractvalue(int64Ty, rangeDescriptor, pos(2)); Value baseSize = baseDesc.size(rank); // Bound upper by base view upper bound. max = llvm_select(llvm_icmp(ICmpPredicate::slt, max, baseSize), max, baseSize); Value size = sub(max, min); // Bound lower by zero. size = llvm_select(llvm_icmp(ICmpPredicate::slt, size, zero), zero, size); Value stride = mul(strides[rank], step); desc.setSize(numNewDims, size); desc.setStride(numNewDims, stride); ++numNewDims; } } rewriter.replaceOp(op, {desc}); return matchSuccess(); } }; /// Conversion pattern that transforms a linalg.transpose op into: /// 1. A function entry `alloca` operation to allocate a ViewDescriptor. /// 2. A load of the ViewDescriptor from the pointer allocated in 1. /// 3. Updates to the ViewDescriptor to introduce the data ptr, offset, size /// and stride. Size and stride are permutations of the original values. /// 4. A store of the resulting ViewDescriptor to the alloca'ed pointer. /// The linalg.transpose op is replaced by the alloca'ed pointer. class TransposeOpConversion : public LLVMOpLowering { public: explicit TransposeOpConversion(MLIRContext *context, LLVMTypeConverter &lowering_) : LLVMOpLowering(TransposeOp::getOperationName(), context, lowering_) {} PatternMatchResult matchAndRewrite(Operation *op, ArrayRef operands, ConversionPatternRewriter &rewriter) const override { // Initialize the common boilerplate and alloca at the top of the FuncOp. edsc::ScopedContext context(rewriter, op->getLoc()); TransposeOpOperandAdaptor adaptor(operands); BaseViewConversionHelper baseDesc(adaptor.view()); auto transposeOp = cast(op); // No permutation, early exit. if (transposeOp.permutation().isIdentity()) return rewriter.replaceOp(op, {baseDesc}), matchSuccess(); BaseViewConversionHelper desc( lowering.convertType(transposeOp.getShapedType())); // Copy the base and aligned pointers from the old descriptor to the new // one. desc.setAllocatedPtr(baseDesc.allocatedPtr()); desc.setAlignedPtr(baseDesc.alignedPtr()); // Copy the offset pointer from the old descriptor to the new one. desc.setOffset(baseDesc.offset()); // Iterate over the dimensions and apply size/stride permutation. for (auto en : llvm::enumerate(transposeOp.permutation().getResults())) { int sourcePos = en.index(); int targetPos = en.value().cast().getPosition(); desc.setSize(targetPos, baseDesc.size(sourcePos)); desc.setStride(targetPos, baseDesc.stride(sourcePos)); } rewriter.replaceOp(op, {desc}); return matchSuccess(); } }; // YieldOp produces and LLVM::ReturnOp. class YieldOpConversion : public LLVMOpLowering { public: explicit YieldOpConversion(MLIRContext *context, LLVMTypeConverter &lowering_) : LLVMOpLowering(YieldOp::getOperationName(), context, lowering_) {} PatternMatchResult matchAndRewrite(Operation *op, ArrayRef operands, ConversionPatternRewriter &rewriter) const override { rewriter.replaceOpWithNewOp(op, operands); return matchSuccess(); } }; template static SmallVector ExtractOperandTypes(Operation *op) { return SmallVector{op->getOperandTypes()}; } template <> SmallVector ExtractOperandTypes(Operation *op) { auto ctx = op->getContext(); auto indexedGenericOp = cast(op); auto numLoops = indexedGenericOp.getNumLoops(); SmallVector result; result.reserve(numLoops + op->getNumOperands()); for (unsigned i = 0; i < numLoops; ++i) { result.push_back(IndexType::get(ctx)); } for (auto type : op->getOperandTypes()) { result.push_back(type); } return result; } // Get a SymbolRefAttr containing the library function name for the LinalgOp. // If the library function does not exist, insert a declaration. template static FlatSymbolRefAttr getLibraryCallSymbolRef(Operation *op, PatternRewriter &rewriter) { auto linalgOp = cast(op); auto fnName = linalgOp.getLibraryCallName(); if (fnName.empty()) { op->emitWarning("No library call defined for: ") << *op; return {}; } // fnName is a dynamic std::String, unique it via a SymbolRefAttr. FlatSymbolRefAttr fnNameAttr = rewriter.getSymbolRefAttr(fnName); auto module = op->getParentOfType(); if (module.lookupSymbol(fnName)) { return fnNameAttr; } SmallVector inputTypes(ExtractOperandTypes(op)); assert(op->getNumResults() == 0 && "Library call for linalg operation can be generated only for ops that " "have void return types"); auto libFnType = FunctionType::get(inputTypes, {}, rewriter.getContext()); OpBuilder::InsertionGuard guard(rewriter); // Insert before module terminator. rewriter.setInsertionPoint(module.getBody(), std::prev(module.getBody()->end())); rewriter.create(op->getLoc(), fnNameAttr.getValue(), libFnType, ArrayRef{}); return fnNameAttr; } } // namespace Type LinalgTypeConverter::convertType(Type t) { if (auto result = LLVMTypeConverter::convertType(t)) return result; return convertLinalgType(t, *this); } namespace { // LinalgOpConversion creates a new call to the // `LinalgOp::getLibraryCallName()` function. // The implementation of the function can be either in the same module or in an // externally linked library. template class LinalgOpConversion : public OpRewritePattern { public: using OpRewritePattern::OpRewritePattern; PatternMatchResult matchAndRewrite(LinalgOp op, PatternRewriter &rewriter) const override { auto libraryCallName = getLibraryCallSymbolRef(op, rewriter); if (!libraryCallName) return this->matchFailure(); rewriter.replaceOpWithNewOp( op, libraryCallName.getValue(), ArrayRef{}, op.getOperands()); return this->matchSuccess(); } }; /// Conversion pattern specialization for CopyOp. This kicks in when both input /// and output permutations are left unspecified or are the identity. template <> class LinalgOpConversion : public OpRewritePattern { public: using OpRewritePattern::OpRewritePattern; PatternMatchResult matchAndRewrite(CopyOp op, PatternRewriter &rewriter) const override { auto inputPerm = op.inputPermutation(); if (inputPerm.hasValue() && !inputPerm->isIdentity()) return matchFailure(); auto outputPerm = op.outputPermutation(); if (outputPerm.hasValue() && !outputPerm->isIdentity()) return matchFailure(); auto libraryCallName = getLibraryCallSymbolRef(op, rewriter); if (!libraryCallName) return matchFailure(); rewriter.replaceOpWithNewOp( op, libraryCallName.getValue(), ArrayRef{}, op.getOperands()); return matchSuccess(); } }; /// Conversion pattern specialization for IndexedGenericOp. template <> class LinalgOpConversion : public OpRewritePattern { public: using OpRewritePattern::OpRewritePattern; PatternMatchResult matchAndRewrite(IndexedGenericOp op, PatternRewriter &rewriter) const override { auto libraryCallName = getLibraryCallSymbolRef(op, rewriter); if (!libraryCallName) return this->matchFailure(); // TODO(pifon, ntv): Use induction variables values instead of zeros, when // IndexedGenericOp is tiled. auto zero = rewriter.create( op.getLoc(), rewriter.getIntegerAttr(rewriter.getIndexType(), 0)); auto indexedGenericOp = cast(op); auto numLoops = indexedGenericOp.getNumLoops(); SmallVector operands; operands.reserve(numLoops + op.getNumOperands()); for (unsigned i = 0; i < numLoops; ++i) { operands.push_back(zero); } for (auto operand : op.getOperands()) { operands.push_back(operand); } rewriter.replaceOpWithNewOp(op, libraryCallName.getValue(), ArrayRef{}, operands); return this->matchSuccess(); } }; /// A non-conversion rewrite pattern kicks in to convert CopyOp with /// permutations into a sequence of TransposeOp and permutation-free CopyOp. /// This interplays together with TransposeOpConversion and /// LinalgConversion to create a path to the LLVM dialect. class CopyTransposeConversion : public OpRewritePattern { public: using OpRewritePattern::OpRewritePattern; PatternMatchResult matchAndRewrite(CopyOp op, PatternRewriter &rewriter) const override { Value in = op.input(), out = op.output(); // If either inputPerm or outputPerm are non-identities, insert transposes. auto inputPerm = op.inputPermutation(); if (inputPerm.hasValue() && !inputPerm->isIdentity()) in = rewriter.create(op.getLoc(), in, AffineMapAttr::get(*inputPerm)); auto outputPerm = op.outputPermutation(); if (outputPerm.hasValue() && !outputPerm->isIdentity()) out = rewriter.create( op.getLoc(), out, AffineMapAttr::get(*outputPerm)); // If nothing was transposed, fail and let the conversion kick in. if (in == op.input() && out == op.output()) return matchFailure(); rewriter.replaceOpWithNewOp(op, in, out); return matchSuccess(); } }; /// Populate the given list with patterns that convert from Linalg to Standard. static void populateLinalgToStandardConversionPatterns(OwningRewritePatternList &patterns, MLIRContext *ctx) { // TODO(ntv) ConvOp conversion needs to export a descriptor with relevant // attribute values such as kernel striding and dilation. patterns.insert, LinalgOpConversion, LinalgOpConversion, LinalgOpConversion, LinalgOpConversion, LinalgOpConversion, LinalgOpConversion, LinalgOpConversion>( ctx); } } // namespace /// Populate the given list with patterns that convert from Linalg to LLVM. void mlir::populateLinalgToLLVMConversionPatterns( LinalgTypeConverter &converter, OwningRewritePatternList &patterns, MLIRContext *ctx) { patterns.insert(ctx, converter); } namespace { struct ConvertLinalgToLLVMPass : public ModulePass { void runOnModule() override; }; } // namespace void ConvertLinalgToLLVMPass::runOnModule() { auto module = getModule(); // Convert to the LLVM IR dialect using the converter defined above. OwningRewritePatternList patterns; LinalgTypeConverter converter(&getContext()); populateAffineToStdConversionPatterns(patterns, &getContext()); populateLoopToStdConversionPatterns(patterns, &getContext()); populateStdToLLVMConversionPatterns(converter, patterns); populateVectorToLLVMConversionPatterns(converter, patterns); populateLinalgToStandardConversionPatterns(patterns, &getContext()); populateLinalgToLLVMConversionPatterns(converter, patterns, &getContext()); ConversionTarget target(getContext()); target.addLegalDialect(); target.addDynamicallyLegalOp( [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); }); target.addLegalOp(); if (failed(applyFullConversion(module, target, patterns, &converter))) signalPassFailure(); } std::unique_ptr> mlir::linalg::createConvertLinalgToLLVMPass() { return std::make_unique(); } static PassRegistration pass( "convert-linalg-to-llvm", "Convert the operations from the linalg dialect into the LLVM dialect");