diff options
author | Florian Mayer <fmayer@google.com> | 2024-02-23 11:31:14 -0800 |
---|---|---|
committer | Florian Mayer <fmayer@google.com> | 2024-02-23 11:31:14 -0800 |
commit | 886b4bc97b0ed5a5e041a0117a584182fc7989c1 (patch) | |
tree | 43cdc0e15e12c298c09251dda38e834e7e778049 /mlir/lib | |
parent | af8afe08ee20a04b2ccb363cac66aa02cfaecd02 (diff) | |
parent | 8d536f83545f071948888983e2db25ce23a8302d (diff) | |
download | llvm-886b4bc97b0ed5a5e041a0117a584182fc7989c1.zip llvm-886b4bc97b0ed5a5e041a0117a584182fc7989c1.tar.gz llvm-886b4bc97b0ed5a5e041a0117a584182fc7989c1.tar.bz2 |
Created using spr 1.3.4
Diffstat (limited to 'mlir/lib')
19 files changed, 929 insertions, 1215 deletions
diff --git a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp index 1b72961..23e9572 100644 --- a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp +++ b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp @@ -148,10 +148,10 @@ struct ExpM1OpLowering : public ConvertOpToLLVMPattern<math::ExpM1Op> { return LLVM::detail::handleMultidimensionalVectors( op.getOperation(), adaptor.getOperands(), *getTypeConverter(), [&](Type llvm1DVectorTy, ValueRange operands) { + auto numElements = LLVM::getVectorNumElements(llvm1DVectorTy); auto splatAttr = SplatElementsAttr::get( - mlir::VectorType::get( - {LLVM::getVectorNumElements(llvm1DVectorTy).getFixedValue()}, - floatType), + mlir::VectorType::get({numElements.getKnownMinValue()}, floatType, + {numElements.isScalable()}), floatOne); auto one = rewriter.create<LLVM::ConstantOp>(loc, llvm1DVectorTy, splatAttr); @@ -207,10 +207,10 @@ struct Log1pOpLowering : public ConvertOpToLLVMPattern<math::Log1pOp> { return LLVM::detail::handleMultidimensionalVectors( op.getOperation(), adaptor.getOperands(), *getTypeConverter(), [&](Type llvm1DVectorTy, ValueRange operands) { + auto numElements = LLVM::getVectorNumElements(llvm1DVectorTy); auto splatAttr = SplatElementsAttr::get( - mlir::VectorType::get( - {LLVM::getVectorNumElements(llvm1DVectorTy).getFixedValue()}, - floatType), + mlir::VectorType::get({numElements.getKnownMinValue()}, floatType, + {numElements.isScalable()}), floatOne); auto one = rewriter.create<LLVM::ConstantOp>(loc, llvm1DVectorTy, splatAttr); @@ -266,10 +266,10 @@ struct RsqrtOpLowering : public ConvertOpToLLVMPattern<math::RsqrtOp> { return LLVM::detail::handleMultidimensionalVectors( op.getOperation(), adaptor.getOperands(), *getTypeConverter(), [&](Type llvm1DVectorTy, ValueRange operands) { + auto numElements = LLVM::getVectorNumElements(llvm1DVectorTy); auto splatAttr = SplatElementsAttr::get( - mlir::VectorType::get( - {LLVM::getVectorNumElements(llvm1DVectorTy).getFixedValue()}, - floatType), + mlir::VectorType::get({numElements.getKnownMinValue()}, floatType, + {numElements.isScalable()}), floatOne); auto one = rewriter.create<LLVM::ConstantOp>(loc, llvm1DVectorTy, splatAttr); diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp index 7eb32eb..7c477f2 100644 --- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp +++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp @@ -384,23 +384,23 @@ createLinalgBodyCalculationForElementwiseOp(Operation *op, ValueRange args, if (isa<tosa::ClampOp>(op) && isa<IntegerType>(elementTy)) { auto intTy = cast<IntegerType>(elementTy); - int32_t min = static_cast<int32_t>( - cast<IntegerAttr>(op->getAttr("min_int")).getValue().getSExtValue()); - int32_t max = static_cast<int32_t>( - cast<IntegerAttr>(op->getAttr("max_int")).getValue().getSExtValue()); + int64_t min = + cast<IntegerAttr>(op->getAttr("min_int")).getValue().getSExtValue(); + int64_t max = + cast<IntegerAttr>(op->getAttr("max_int")).getValue().getSExtValue(); if (intTy.isUnsignedInteger()) { - min = std::max<int32_t>(min, 0); - max = std::min<int32_t>( + min = std::max(min, (int64_t)0); + max = std::min( max, APInt::getMaxValue(intTy.getIntOrFloatBitWidth()).getSExtValue()); } else { - min = std::max<int32_t>( - min, APInt::getSignedMinValue(intTy.getIntOrFloatBitWidth()) - .getSExtValue()); - max = std::min<int32_t>( - max, APInt::getSignedMaxValue(intTy.getIntOrFloatBitWidth()) - .getSExtValue()); + min = + std::max(min, APInt::getSignedMinValue(intTy.getIntOrFloatBitWidth()) + .getSExtValue()); + max = + std::min(max, APInt::getSignedMaxValue(intTy.getIntOrFloatBitWidth()) + .getSExtValue()); } auto minVal = rewriter.create<arith::ConstantIntOp>( diff --git a/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp b/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp index e88f82c..26dfb38 100644 --- a/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp +++ b/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp @@ -40,12 +40,12 @@ namespace { //===----------------------------------------------------------------------===// // Common match failure reasons. -static constexpr StringLiteral MATCH_FAILURE_NOT_SME_TILE_TYPE_MULTIPLE( +static constexpr StringLiteral kMatchFailureNotSMETileTypeMultiple( "op vector size is not multiple of SME tiles"); -static constexpr StringLiteral MATCH_FAILURE_UNSUPPORTED_MASK_OP( +static constexpr StringLiteral kMatchFailureUnsupportedMaskOp( "op mask is unsupported for legalization/decomposition"); static constexpr StringLiteral - MATCH_FAILURE_NON_PERMUTATION_MAP("op affine map is not a permutation"); + kMatchFailureNonPermutationMap("op affine map is not a permutation"); /// An SMESubTile represents a single SME-sized sub-tile from decomposing a /// larger vector type. The (`row`, `col`) are the position of the tile in the @@ -174,8 +174,8 @@ struct LegalizeVectorOuterProductOpsByDecomposition OneToNPatternRewriter &rewriter) const override { auto vectorType = outerProductOp.getResultVectorType(); if (!isMultipleOfSMETileVectorType(vectorType)) - return rewriter.notifyMatchFailure( - outerProductOp, MATCH_FAILURE_NOT_SME_TILE_TYPE_MULTIPLE); + return rewriter.notifyMatchFailure(outerProductOp, + kMatchFailureNotSMETileTypeMultiple); Value mask; Operation *rootOp = outerProductOp; @@ -188,7 +188,7 @@ struct LegalizeVectorOuterProductOpsByDecomposition if (!isSupportedMaskOp(mask)) return rewriter.notifyMatchFailure(outerProductOp, - MATCH_FAILURE_UNSUPPORTED_MASK_OP); + kMatchFailureUnsupportedMaskOp); ValueRange accSMETiles = adaptor.getAcc(); auto smeTileType = getSMETileTypeForElement(vectorType.getElementType()); @@ -252,18 +252,18 @@ struct LegalizeTransferReadOpsByDecomposition OneToNPatternRewriter &rewriter) const override { auto vectorType = readOp.getVectorType(); if (!isMultipleOfSMETileVectorType(vectorType)) - return rewriter.notifyMatchFailure( - readOp, MATCH_FAILURE_NOT_SME_TILE_TYPE_MULTIPLE); + return rewriter.notifyMatchFailure(readOp, + kMatchFailureNotSMETileTypeMultiple); auto mask = readOp.getMask(); if (!isSupportedMaskOp(mask)) return rewriter.notifyMatchFailure(readOp, - MATCH_FAILURE_UNSUPPORTED_MASK_OP); + kMatchFailureUnsupportedMaskOp); auto permutationMap = readOp.getPermutationMap(); if (!permutationMap.isPermutation()) return rewriter.notifyMatchFailure(readOp, - MATCH_FAILURE_NON_PERMUTATION_MAP); + kMatchFailureNonPermutationMap); // Note: For 2D vector types the only non-identity permutation is a simple // tranpose [1, 0]. @@ -300,18 +300,18 @@ struct LegalizeTransferWriteOpsByDecomposition OneToNPatternRewriter &rewriter) const override { auto vectorType = writeOp.getVectorType(); if (!isMultipleOfSMETileVectorType(vectorType)) - return rewriter.notifyMatchFailure( - writeOp, MATCH_FAILURE_NOT_SME_TILE_TYPE_MULTIPLE); + return rewriter.notifyMatchFailure(writeOp, + kMatchFailureNotSMETileTypeMultiple); auto mask = writeOp.getMask(); if (!isSupportedMaskOp(mask)) return rewriter.notifyMatchFailure(writeOp, - MATCH_FAILURE_UNSUPPORTED_MASK_OP); + kMatchFailureUnsupportedMaskOp); auto permutationMap = writeOp.getPermutationMap(); if (!permutationMap.isPermutation()) return rewriter.notifyMatchFailure(writeOp, - MATCH_FAILURE_NON_PERMUTATION_MAP); + kMatchFailureNonPermutationMap); // Note: For 2D vector types the only non-identity permutation is a simple // tranpose [1, 0]. diff --git a/mlir/lib/Dialect/GPU/CMakeLists.txt b/mlir/lib/Dialect/GPU/CMakeLists.txt index e5776e1..51cfa22 100644 --- a/mlir/lib/Dialect/GPU/CMakeLists.txt +++ b/mlir/lib/Dialect/GPU/CMakeLists.txt @@ -1,11 +1,3 @@ -if ("NVPTX" IN_LIST LLVM_TARGETS_TO_BUILD) - set(NVPTX_LIBS - NVPTXCodeGen - NVPTXDesc - NVPTXInfo - ) -endif() - if (MLIR_ENABLE_ROCM_CONVERSIONS) set(AMDGPU_LIBS IRReader @@ -60,7 +52,6 @@ add_mlir_dialect_library(MLIRGPUTransforms Transforms/ParallelLoopMapper.cpp Transforms/ROCDLAttachTarget.cpp Transforms/SerializeToBlob.cpp - Transforms/SerializeToCubin.cpp Transforms/SerializeToHsaco.cpp Transforms/ShuffleRewriter.cpp Transforms/SPIRVAttachTarget.cpp @@ -74,7 +65,6 @@ add_mlir_dialect_library(MLIRGPUTransforms Core MC Target - ${NVPTX_LIBS} ${AMDGPU_LIBS} DEPENDS @@ -110,48 +100,6 @@ add_mlir_dialect_library(MLIRGPUTransforms add_subdirectory(TransformOps) add_subdirectory(Pipelines) -if(MLIR_ENABLE_CUDA_RUNNER) - if(NOT MLIR_ENABLE_CUDA_CONVERSIONS) - message(SEND_ERROR - "Building mlir with cuda support requires the NVPTX backend") - endif() - - # Configure CUDA language support. Using check_language first allows us to - # give a custom error message. - include(CheckLanguage) - check_language(CUDA) - if (CMAKE_CUDA_COMPILER) - enable_language(CUDA) - else() - message(SEND_ERROR - "Building mlir with cuda support requires a working CUDA install") - endif() - - # Enable gpu-to-cubin pass. - target_compile_definitions(obj.MLIRGPUTransforms - PRIVATE - MLIR_GPU_TO_CUBIN_PASS_ENABLE=1 - ) - - # Add CUDA headers includes and the libcuda.so library. - target_include_directories(obj.MLIRGPUTransforms - PRIVATE - ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} - ) - - # Add link path for the cuda driver library. - find_library(CUDA_DRIVER_LIBRARY cuda HINTS ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES} REQUIRED) - get_filename_component(CUDA_DRIVER_LIBRARY_PATH "${CUDA_DRIVER_LIBRARY}" DIRECTORY) - target_link_directories(MLIRGPUTransforms PRIVATE ${CUDA_DRIVER_LIBRARY_PATH}) - - target_link_libraries(MLIRGPUTransforms - PRIVATE - MLIRNVVMToLLVMIRTranslation - cuda - ) - -endif() - if(MLIR_ENABLE_ROCM_CONVERSIONS) if (NOT ("AMDGPU" IN_LIST LLVM_TARGETS_TO_BUILD)) message(SEND_ERROR diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp index 30b6cd7..33ce5c1 100644 --- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp +++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp @@ -648,6 +648,8 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result, TypeRange workgroupAttributions, TypeRange privateAttributions, Value clusterSizeX, Value clusterSizeY, Value clusterSizeZ) { + OpBuilder::InsertionGuard g(builder); + // Add a WorkGroup attribution attribute. This attribute is required to // identify private attributions in the list of block argguments. result.addAttribute(getNumWorkgroupAttributionsAttrName(), @@ -674,7 +676,7 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result, // attributions, where the first kNumConfigRegionAttributes arguments have // `index` type and the rest have the same types as the data operands. Region *kernelRegion = result.addRegion(); - Block *body = new Block(); + Block *body = builder.createBlock(kernelRegion); // TODO: Allow passing in proper locations here. for (unsigned i = 0; i < kNumConfigRegionAttributes; ++i) body->addArgument(builder.getIndexType(), result.location); @@ -683,7 +685,6 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result, body->addArgument(argTy, result.location); for (Type argTy : privateAttributions) body->addArgument(argTy, result.location); - kernelRegion->push_back(body); // Fill OperandSegmentSize Attribute. SmallVector<int32_t, 11> segmentSizes(11, 1); segmentSizes.front() = asyncDependencies.size(); @@ -1325,6 +1326,8 @@ void GPUFuncOp::build(OpBuilder &builder, OperationState &result, TypeRange workgroupAttributions, TypeRange privateAttributions, ArrayRef<NamedAttribute> attrs) { + OpBuilder::InsertionGuard g(builder); + result.addAttribute(SymbolTable::getSymbolAttrName(), builder.getStringAttr(name)); result.addAttribute(getFunctionTypeAttrName(result.name), @@ -1333,7 +1336,7 @@ void GPUFuncOp::build(OpBuilder &builder, OperationState &result, builder.getI64IntegerAttr(workgroupAttributions.size())); result.addAttributes(attrs); Region *body = result.addRegion(); - Block *entryBlock = new Block; + Block *entryBlock = builder.createBlock(body); // TODO: Allow passing in proper locations here. for (Type argTy : type.getInputs()) @@ -1342,8 +1345,6 @@ void GPUFuncOp::build(OpBuilder &builder, OperationState &result, entryBlock->addArgument(argTy, result.location); for (Type argTy : privateAttributions) entryBlock->addArgument(argTy, result.location); - - body->getBlocks().push_back(entryBlock); } /// Parses a GPU function memory attribution. diff --git a/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp b/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp deleted file mode 100644 index 34ad4e6..0000000 --- a/mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp +++ /dev/null @@ -1,180 +0,0 @@ -//===- LowerGPUToCUBIN.cpp - Convert GPU kernel to CUBIN blob -------------===// -// -// 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 implements a pass that serializes a gpu module into CUBIN blob and -// adds that blob as a string attribute of the module. -// -//===----------------------------------------------------------------------===// - -#include "mlir/Dialect/GPU/Transforms/Passes.h" -#include "mlir/Dialect/LLVMIR/NVVMDialect.h" -#include "llvm/Support/Debug.h" - -#if MLIR_GPU_TO_CUBIN_PASS_ENABLE -#include "mlir/Pass/Pass.h" -#include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h" -#include "mlir/Target/LLVMIR/Export.h" -#include "llvm/Support/TargetSelect.h" -#include "llvm/Support/Threading.h" - -#include <cuda.h> - -using namespace mlir; - -static void emitCudaError(const llvm::Twine &expr, const char *buffer, - CUresult result, Location loc) { - const char *error = nullptr; - cuGetErrorString(result, &error); - emitError(loc, - expr.concat(error ? " failed with error code " + llvm::Twine{error} - : llvm::Twine(" failed with unknown error ")) - .concat("[") - .concat(buffer) - .concat("]")); -} - -#define RETURN_ON_CUDA_ERROR(expr) \ - do { \ - if (auto status = (expr)) { \ - emitCudaError(#expr, jitErrorBuffer, status, loc); \ - return {}; \ - } \ - } while (false) - -namespace { -class SerializeToCubinPass - : public PassWrapper<SerializeToCubinPass, gpu::SerializeToBlobPass> { - static llvm::once_flag initializeBackendOnce; - -public: - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SerializeToCubinPass) - - SerializeToCubinPass(StringRef triple = "nvptx64-nvidia-cuda", - StringRef chip = "sm_35", StringRef features = "+ptx60", - int optLevel = 2, bool dumpPtx = false); - - StringRef getArgument() const override { return "gpu-to-cubin"; } - StringRef getDescription() const override { - return "Lower GPU kernel function to CUBIN binary annotations"; - } - -private: - // Serializes PTX to CUBIN. - std::unique_ptr<std::vector<char>> - serializeISA(const std::string &isa) override; -}; -} // namespace - -// Sets the 'option' to 'value' unless it already has a value. -static void maybeSetOption(Pass::Option<std::string> &option, StringRef value) { - if (!option.hasValue()) - option = value.str(); -} - -llvm::once_flag SerializeToCubinPass::initializeBackendOnce; - -SerializeToCubinPass::SerializeToCubinPass(StringRef triple, StringRef chip, - StringRef features, int optLevel, - bool dumpPtx) { - // No matter how this pass is constructed, ensure that the NVPTX backend - // is initialized exactly once. - llvm::call_once(initializeBackendOnce, []() { - // Initialize LLVM NVPTX backend. -#if LLVM_HAS_NVPTX_TARGET - LLVMInitializeNVPTXTarget(); - LLVMInitializeNVPTXTargetInfo(); - LLVMInitializeNVPTXTargetMC(); - LLVMInitializeNVPTXAsmPrinter(); -#endif - }); - - maybeSetOption(this->triple, triple); - maybeSetOption(this->chip, chip); - maybeSetOption(this->features, features); - this->dumpPtx = dumpPtx; - if (this->optLevel.getNumOccurrences() == 0) - this->optLevel.setValue(optLevel); -} - -std::unique_ptr<std::vector<char>> -SerializeToCubinPass::serializeISA(const std::string &isa) { - Location loc = getOperation().getLoc(); - char jitErrorBuffer[4096] = {0}; - - RETURN_ON_CUDA_ERROR(cuInit(0)); - - // Linking requires a device context. - CUdevice device; - RETURN_ON_CUDA_ERROR(cuDeviceGet(&device, 0)); - CUcontext context; - // Use the primary context. - RETURN_ON_CUDA_ERROR(cuDevicePrimaryCtxRetain(&context, device)); - // Push the primary context so that the next CUDA operations - // actually use it. - RETURN_ON_CUDA_ERROR(cuCtxPushCurrent(context)); - CUlinkState linkState; - - CUjit_option jitOptions[] = {CU_JIT_ERROR_LOG_BUFFER, - CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES}; - void *jitOptionsVals[] = {jitErrorBuffer, - reinterpret_cast<void *>(sizeof(jitErrorBuffer))}; - - RETURN_ON_CUDA_ERROR(cuLinkCreate(2, /* number of jit options */ - jitOptions, /* jit options */ - jitOptionsVals, /* jit option values */ - &linkState)); - - auto kernelName = getOperation().getName().str(); - if (dumpPtx) { - llvm::dbgs() << " Kernel Name : [" << kernelName << "]\n"; - llvm::dbgs() << isa << "\n"; - } - RETURN_ON_CUDA_ERROR(cuLinkAddData( - linkState, CUjitInputType::CU_JIT_INPUT_PTX, - const_cast<void *>(static_cast<const void *>(isa.c_str())), isa.length(), - kernelName.c_str(), 0, /* number of jit options */ - nullptr, /* jit options */ - nullptr /* jit option values */ - )); - - void *cubinData; - size_t cubinSize; - RETURN_ON_CUDA_ERROR(cuLinkComplete(linkState, &cubinData, &cubinSize)); - - char *cubinAsChar = static_cast<char *>(cubinData); - auto result = - std::make_unique<std::vector<char>>(cubinAsChar, cubinAsChar + cubinSize); - - // This will also destroy the cubin data. - RETURN_ON_CUDA_ERROR(cuLinkDestroy(linkState)); - // Pop and release the primary context. - CUcontext poppedContext; - RETURN_ON_CUDA_ERROR(cuCtxPopCurrent(&poppedContext)); - RETURN_ON_CUDA_ERROR(cuDevicePrimaryCtxRelease(device)); - - return result; -} - -// Register pass to serialize GPU kernel functions to a CUBIN binary annotation. -void mlir::registerGpuSerializeToCubinPass() { - PassRegistration<SerializeToCubinPass> registerSerializeToCubin( - [] { return std::make_unique<SerializeToCubinPass>(); }); -} - -std::unique_ptr<Pass> mlir::createGpuSerializeToCubinPass(StringRef triple, - StringRef arch, - StringRef features, - int optLevel, - bool dumpPtx) { - return std::make_unique<SerializeToCubinPass>(triple, arch, features, - optLevel, dumpPtx); -} - -#else // MLIR_GPU_TO_CUBIN_PASS_ENABLE -void mlir::registerGpuSerializeToCubinPass() {} -#endif // MLIR_GPU_TO_CUBIN_PASS_ENABLE diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp index 7eed792..3627ff6 100644 --- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp +++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp @@ -1041,6 +1041,11 @@ int64_t LinalgOp::getIndexingMapIndex(OpOperand *opOperand) { LogicalResult mlir::linalg::detail::verifyStructuredOpInterface(Operation *op) { LinalgOp linalgOp = cast<LinalgOp>(op); + // Mixed tensor/buffer operands are not allowed. + if (!linalgOp.hasPureTensorSemantics() && + !linalgOp.hasPureBufferSemantics() && op->getNumOperands() > 0) + return op->emitOpError("expected to have pure tensor or buffer semantics"); + // Before checking indexing maps, we need to make sure the attributes // referenced by it are valid. if (linalgOp.hasDynamicIndexingMaps()) diff --git a/mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp b/mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp index 71e4e13..962cb28 100644 --- a/mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp +++ b/mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp @@ -1471,6 +1471,16 @@ RsqrtApproximation::matchAndRewrite(math::RsqrtOp op, //----------------------------------------------------------------------------// +void mlir::populatePolynomialApproximateTanhPattern( + RewritePatternSet &patterns) { + patterns.add<TanhApproximation>(patterns.getContext()); +} + +void mlir::populatePolynomialApproximateErfPattern( + RewritePatternSet &patterns) { + patterns.add<ErfPolynomialApproximation>(patterns.getContext()); +} + void mlir::populateMathPolynomialApproximationPatterns( RewritePatternSet &patterns, const MathPolynomialApproximationOptions &options) { diff --git a/mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp b/mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp index 7cbe0de..c4d8b0b 100644 --- a/mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp +++ b/mlir/lib/Dialect/Mesh/Transforms/Spmdization.cpp @@ -593,7 +593,6 @@ static SmallVector<MeshShardingAttr> getOperandShardings(Operation &op) { Operation *definingOp = operand.getDefiningOp(); assert(definingOp); ShardOp shardOp = llvm::cast<ShardOp>(definingOp); - assert(shardOp.getAnnotateForUsers()); return shardOp.getShard(); }); return res; @@ -615,34 +614,46 @@ static SmallVector<MeshShardingAttr> getResultShardings(Operation &op) { assert(result.hasOneUse()); Operation *userOp = *result.getUsers().begin(); ShardOp shardOp = llvm::cast<ShardOp>(userOp); - assert(!shardOp.getAnnotateForUsers()); return shardOp.getShard(); }); return res; } static LogicalResult -spmdizeOperation(Operation &op, IRMapping &spmdizationMap, +spmdizeOperation(ShardOp shardOp, IRMapping &spmdizationMap, SymbolTableCollection &symbolTableCollection, OpBuilder &builder) { - ShardOp shardOp = llvm::dyn_cast<ShardOp>(op); - if (shardOp) { - if (!shardOp.getAnnotateForUsers()) { - return success(); - } - + Value targetSpmdValue; + + // Check if 2 shard ops are chained. If not there is no need for resharding + // as the source and target shared the same sharding. + ShardOp srcShardOp = + dyn_cast_or_null<ShardOp>(shardOp.getOperand().getDefiningOp()); + if (!srcShardOp) { + targetSpmdValue = spmdizationMap.lookup(shardOp.getOperand()); + } else { // Insert resharding. - ShardOp srcShardOp = - llvm::cast<ShardOp>(shardOp.getOperand().getDefiningOp()); - assert(!srcShardOp.getAnnotateForUsers()); + assert(!srcShardOp.getAnnotateForUsers() && shardOp.getAnnotateForUsers()); TypedValue<ShapedType> srcSpmdValue = spmdizationMap.lookup(srcShardOp.getOperand()) .cast<TypedValue<ShapedType>>(); - Value targetSpmdValue = reshard(builder, srcShardOp, shardOp, srcSpmdValue, - symbolTableCollection); - assert(!spmdizationMap.contains(shardOp.getResult())); - spmdizationMap.map(shardOp.getResult(), targetSpmdValue); - return success(); + targetSpmdValue = reshard(builder, srcShardOp, shardOp, srcSpmdValue, + symbolTableCollection); + } + + assert(!spmdizationMap.contains(shardOp.getResult())); + spmdizationMap.map(shardOp.getResult(), targetSpmdValue); + return success(); +} + +static LogicalResult +spmdizeOperation(Operation &op, IRMapping &spmdizationMap, + SymbolTableCollection &symbolTableCollection, + OpBuilder &builder) { + ShardOp shardOp = llvm::dyn_cast<ShardOp>(op); + if (shardOp) { + return spmdizeOperation(shardOp, spmdizationMap, symbolTableCollection, + builder); } SmallVector<Value> spmdizedOperands; diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp index 119df9a..233e702 100644 --- a/mlir/lib/Dialect/SCF/IR/SCF.cpp +++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp @@ -306,17 +306,18 @@ void ConditionOp::getSuccessorRegions( void ForOp::build(OpBuilder &builder, OperationState &result, Value lb, Value ub, Value step, ValueRange iterArgs, BodyBuilderFn bodyBuilder) { + OpBuilder::InsertionGuard guard(builder); + result.addOperands({lb, ub, step}); result.addOperands(iterArgs); for (Value v : iterArgs) result.addTypes(v.getType()); Type t = lb.getType(); Region *bodyRegion = result.addRegion(); - bodyRegion->push_back(new Block); - Block &bodyBlock = bodyRegion->front(); - bodyBlock.addArgument(t, result.location); + Block *bodyBlock = builder.createBlock(bodyRegion); + bodyBlock->addArgument(t, result.location); for (Value v : iterArgs) - bodyBlock.addArgument(v.getType(), v.getLoc()); + bodyBlock->addArgument(v.getType(), v.getLoc()); // Create the default terminator if the builder is not provided and if the // iteration arguments are not provided. Otherwise, leave this to the caller @@ -325,9 +326,9 @@ void ForOp::build(OpBuilder &builder, OperationState &result, Value lb, ForOp::ensureTerminator(*bodyRegion, builder, result.location); } else if (bodyBuilder) { OpBuilder::InsertionGuard guard(builder); - builder.setInsertionPointToStart(&bodyBlock); - bodyBuilder(builder, result.location, bodyBlock.getArgument(0), - bodyBlock.getArguments().drop_front()); + builder.setInsertionPointToStart(bodyBlock); + bodyBuilder(builder, result.location, bodyBlock->getArgument(0), + bodyBlock->getArguments().drop_front()); } } diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp index 950ee59..62d0785 100644 --- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp +++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp @@ -1330,6 +1330,7 @@ NARY_SHAPE_INFER(tosa::CastOp) NARY_SHAPE_INFER(tosa::CeilOp) NARY_SHAPE_INFER(tosa::ClampOp) NARY_SHAPE_INFER(tosa::ClzOp) +NARY_SHAPE_INFER(tosa::CosOp) NARY_SHAPE_INFER(tosa::DivOp) NARY_SHAPE_INFER(tosa::ExpOp) NARY_SHAPE_INFER(tosa::FloorOp) @@ -1352,6 +1353,7 @@ NARY_SHAPE_INFER(tosa::ReciprocalOp) NARY_SHAPE_INFER(tosa::RescaleOp) NARY_SHAPE_INFER(tosa::ReverseOp) NARY_SHAPE_INFER(tosa::RsqrtOp) +NARY_SHAPE_INFER(tosa::SinOp) NARY_SHAPE_INFER(tosa::SelectOp) NARY_SHAPE_INFER(tosa::SubOp) NARY_SHAPE_INFER(tosa::TanhOp) diff --git a/mlir/lib/Dialect/Utils/IndexingUtils.cpp b/mlir/lib/Dialect/Utils/IndexingUtils.cpp index baaa581..4c96065 100644 --- a/mlir/lib/Dialect/Utils/IndexingUtils.cpp +++ b/mlir/lib/Dialect/Utils/IndexingUtils.cpp @@ -7,13 +7,12 @@ //===----------------------------------------------------------------------===// #include "mlir/Dialect/Utils/IndexingUtils.h" - +#include "mlir/Dialect/Utils/StaticValueUtils.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/MLIRContext.h" #include "llvm/ADT/STLExtras.h" - #include <numeric> #include <optional> @@ -306,6 +305,14 @@ mlir::computeLinearIndex(OpFoldResult sourceOffset, return {expr, values}; } +std::pair<AffineExpr, SmallVector<OpFoldResult>> +mlir::computeLinearIndex(OpFoldResult sourceOffset, ArrayRef<int64_t> strides, + ArrayRef<Value> indices) { + return computeLinearIndex( + sourceOffset, getAsIndexOpFoldResult(sourceOffset.getContext(), strides), + getAsOpFoldResult(ValueRange(indices))); +} + //===----------------------------------------------------------------------===// // TileOffsetRange //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp index 36fb667..fc11ae6 100644 --- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp +++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp @@ -724,9 +724,8 @@ BitCastRewriter::BitCastRewriter(VectorType sourceVectorType, static LogicalResult commonConversionPrecondition(PatternRewriter &rewriter, VectorType preconditionType, Operation *op) { - if (!preconditionType || preconditionType.getRank() != 1 || - preconditionType.isScalable()) - return rewriter.notifyMatchFailure(op, "scalable or >1-D vector"); + if (!preconditionType || preconditionType.isScalable()) + return rewriter.notifyMatchFailure(op, "scalable vector"); // TODO: consider relaxing this restriction in the future if we find ways // to really work with subbyte elements across the MLIR/LLVM boundary. @@ -743,6 +742,9 @@ LogicalResult BitCastRewriter::commonPrecondition(PatternRewriter &rewriter, if (!enumerator.sourceVectorType || !enumerator.targetVectorType) return rewriter.notifyMatchFailure(op, "types are not vector"); + if (!preconditionType || preconditionType.getRank() != 1) + return rewriter.notifyMatchFailure(op, "unsupported >1-D vector"); + return commonConversionPrecondition(rewriter, preconditionType, op); } @@ -855,7 +857,6 @@ static Value rewriteI4ToI8SignedExt(PatternRewriter &rewriter, Location loc, "Expected i4 type"); // 1. Generate a bitcast vector<Xxi4> -> vector<X/2xi8>. - int64_t vecDimSize = srcVecType.getShape().back(); SmallVector<int64_t> i8VecShape = llvm::to_vector(srcVecType.getShape()); constexpr int64_t i4Toi8BitwidthFactor = 2; i8VecShape.back() = i8VecShape.back() / i4Toi8BitwidthFactor; @@ -871,16 +872,8 @@ static Value rewriteI4ToI8SignedExt(PatternRewriter &rewriter, Location loc, Value low = rewriter.create<arith::ShRSIOp>(loc, shl, shiftValues); Value high = rewriter.create<arith::ShRSIOp>(loc, i8Vector, shiftValues); - // 3. Interleave low and high i8 elements using a shuffle. - SmallVector<int64_t> interleaveMaskValues; - interleaveMaskValues.reserve(vecDimSize); - for (int i = 0, end = vecDimSize / 2; i < end; ++i) { - interleaveMaskValues.push_back(i); - interleaveMaskValues.push_back(i + (vecDimSize / 2)); - } - - return rewriter.create<vector::ShuffleOp>( - loc, low, high, rewriter.getI64ArrayAttr(interleaveMaskValues)); + // 3. Interleave low and high i8 elements. + return rewriter.create<vector::InterleaveOp>(loc, low, high); } namespace { @@ -1008,8 +1001,7 @@ struct RewriteExtOfBitCast : OpRewritePattern<ExtOpType> { /// %1 = arith.shli %0, 4 : vector<4xi8> /// %2 = arith.shrsi %1, 4 : vector<4xi8> /// %3 = arith.shrsi %0, 4 : vector<4xi8> -/// %4 = vector.shuffle %2, %3 [0, 4, 1, 5, 2, 6, 3, 7] -/// : vector<4xi8>, vector<4xi8> +/// %4 = vector.interleave %2, %3 : vector<4xi8> /// %5 = arith.extsi %4 : vector<8xi8> to vector<8xi32> /// /// arith.sitofp %in : vector<8xi4> to vector<8xf32> @@ -1018,8 +1010,7 @@ struct RewriteExtOfBitCast : OpRewritePattern<ExtOpType> { /// %1 = arith.shli %0, 4 : vector<4xi8> /// %2 = arith.shrsi %1, 4 : vector<4xi8> /// %3 = arith.shrsi %0, 4 : vector<4xi8> -/// %4 = vector.shuffle %2, %3 [0, 4, 1, 5, 2, 6, 3, 7] -/// : vector<4xi8>, vector<4xi8> +/// %4 = vector.interleave %2, %3 : vector<4xi8> /// %5 = arith.sitofp %4 : vector<8xi8> to vector<8xf32> /// template <typename ConversionOpType> diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp index 04e5a81..0ffef6a 100644 --- a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp +++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp @@ -15,6 +15,7 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Utils/IndexingUtils.h" #include "mlir/Dialect/Vector/IR/VectorOps.h" #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h" #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h" @@ -577,7 +578,6 @@ public: if (transferReadOp.getMask()) return failure(); - SmallVector<Value> collapsedIndices; int64_t firstDimToCollapse = sourceType.getRank() - vectorType.getRank(); // 1. Collapse the source memref @@ -599,12 +599,14 @@ public: // 2.2 New indices // If all the collapsed indices are zero then no extra logic is needed. // Otherwise, a new offset/index has to be computed. + SmallVector<Value> collapsedIndices; if (failed(checkAndCollapseInnerZeroIndices(transferReadOp.getIndices(), firstDimToCollapse, collapsedIndices))) { - // Copy all the leading indices - collapsedIndices = transferReadOp.getIndices(); - collapsedIndices.resize(firstDimToCollapse); + // Copy all the leading indices. + SmallVector<Value> indices = transferReadOp.getIndices(); + collapsedIndices.append(indices.begin(), + indices.begin() + firstDimToCollapse); // Compute the remaining trailing index/offset required for reading from // the collapsed memref: @@ -621,24 +623,26 @@ public: // memref<1x86xi32>, vector<2xi32> // one would get the following offset: // %offset = %arg0 * 43 - AffineExpr offsetExpr, idxExpr; - bindSymbols(rewriter.getContext(), offsetExpr, idxExpr); - - int64_t outputRank = transferReadOp.getIndices().size(); - OpFoldResult offset = + OpFoldResult collapsedOffset = rewriter.create<arith::ConstantIndexOp>(loc, 0).getResult(); - for (int64_t i = firstDimToCollapse; i < outputRank; ++i) { - int64_t dim = dyn_cast<ShapedType>(source.getType()).getDimSize(i); - offset = affine::makeComposedFoldedAffineApply( - rewriter, loc, offsetExpr + dim * idxExpr, - {offset, transferReadOp.getIndices()[i]}); - } - if (offset.is<Value>()) { - collapsedIndices.push_back(offset.get<Value>()); + auto sourceShape = sourceType.getShape(); + auto collapsedStrides = computeSuffixProduct(ArrayRef<int64_t>( + sourceShape.begin() + firstDimToCollapse, sourceShape.end())); + + // Compute the collapsed offset. + ArrayRef<Value> indicesToCollapse(indices.begin() + firstDimToCollapse, + indices.end()); + auto &&[collapsedExpr, collapsedVals] = computeLinearIndex( + collapsedOffset, collapsedStrides, indicesToCollapse); + collapsedOffset = affine::makeComposedFoldedAffineApply( + rewriter, loc, collapsedExpr, collapsedVals); + + if (collapsedOffset.is<Value>()) { + collapsedIndices.push_back(collapsedOffset.get<Value>()); } else { collapsedIndices.push_back(rewriter.create<arith::ConstantIndexOp>( - loc, *getConstantIntValue(offset))); + loc, *getConstantIntValue(collapsedOffset))); } } @@ -710,6 +714,7 @@ public: firstContiguousInnerDim, collapsedIndices))) return failure(); + Value collapsedSource = collapseInnerDims(rewriter, loc, source, firstContiguousInnerDim); MemRefType collapsedSourceType = diff --git a/mlir/lib/ExecutionEngine/SparseTensor/Storage.cpp b/mlir/lib/ExecutionEngine/SparseTensor/Storage.cpp index 9e8b240..bbe10b0 100644 --- a/mlir/lib/ExecutionEngine/SparseTensor/Storage.cpp +++ b/mlir/lib/ExecutionEngine/SparseTensor/Storage.cpp @@ -74,13 +74,6 @@ MLIR_SPARSETENSOR_FOREVERY_FIXED_O(IMPL_GETCOORDINATES) MLIR_SPARSETENSOR_FOREVERY_V(IMPL_GETVALUES) #undef IMPL_GETVALUES -#define IMPL_FORWARDINGINSERT(VNAME, V) \ - void SparseTensorStorageBase::forwardingInsert(const uint64_t *, V) { \ - FATAL_PIV("forwardingInsert" #VNAME); \ - } -MLIR_SPARSETENSOR_FOREVERY_V(IMPL_FORWARDINGINSERT) -#undef IMPL_FORWARDINGINSERT - #define IMPL_LEXINSERT(VNAME, V) \ void SparseTensorStorageBase::lexInsert(const uint64_t *, V) { \ FATAL_PIV("lexInsert" #VNAME); \ diff --git a/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp b/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp index a5e75a7..0bc90b2 100644 --- a/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp +++ b/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp @@ -117,20 +117,7 @@ extern "C" { switch (action) { \ case Action::kEmpty: { \ return SparseTensorStorage<P, C, V>::newEmpty( \ - dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim, \ - false); \ - } \ - case Action::kEmptyForward: { \ - return SparseTensorStorage<P, C, V>::newEmpty( \ - dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim, \ - true); \ - } \ - case Action::kFromCOO: { \ - assert(ptr && "Received nullptr for SparseTensorCOO object"); \ - auto &coo = *static_cast<SparseTensorCOO<V> *>(ptr); \ - return SparseTensorStorage<P, C, V>::newFromCOO( \ - dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim, \ - coo); \ + dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim); \ } \ case Action::kFromReader: { \ assert(ptr && "Received nullptr for SparseTensorReader object"); \ @@ -138,11 +125,6 @@ extern "C" { return static_cast<void *>(reader.readSparseTensor<P, C, V>( \ lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim)); \ } \ - case Action::kToCOO: { \ - assert(ptr && "Received nullptr for SparseTensorStorage object"); \ - auto &tensor = *static_cast<SparseTensorStorage<P, C, V> *>(ptr); \ - return tensor.toCOO(); \ - } \ case Action::kPack: { \ assert(ptr && "Received nullptr for SparseTensorStorage object"); \ intptr_t *buffers = static_cast<intptr_t *>(ptr); \ @@ -341,21 +323,6 @@ MLIR_SPARSETENSOR_FOREVERY_O(IMPL_SPARSECOORDINATES) #undef IMPL_SPARSECOORDINATES #undef IMPL_GETOVERHEAD -#define IMPL_FORWARDINGINSERT(VNAME, V) \ - void _mlir_ciface_forwardingInsert##VNAME( \ - void *t, StridedMemRefType<V, 0> *vref, \ - StridedMemRefType<index_type, 1> *dimCoordsRef) { \ - assert(t &&vref); \ - ASSERT_NO_STRIDE(dimCoordsRef); \ - const index_type *dimCoords = MEMREF_GET_PAYLOAD(dimCoordsRef); \ - assert(dimCoords); \ - const V *value = MEMREF_GET_PAYLOAD(vref); \ - static_cast<SparseTensorStorageBase *>(t)->forwardingInsert(dimCoords, \ - *value); \ - } -MLIR_SPARSETENSOR_FOREVERY_V(IMPL_FORWARDINGINSERT) -#undef IMPL_FORWARDINGINSERT - #define IMPL_LEXINSERT(VNAME, V) \ void _mlir_ciface_lexInsert##VNAME( \ void *t, StridedMemRefType<index_type, 1> *lvlCoordsRef, \ @@ -427,8 +394,8 @@ void _mlir_ciface_getSparseTensorReaderDimSizes( const uint64_t cSize = MEMREF_GET_USIZE(cref); \ const uint64_t vSize = MEMREF_GET_USIZE(vref); \ ASSERT_USIZE_EQ(lvl2dimRef, dimRank); \ - assert(cSize >= lvlRank * vSize); \ - assert(vSize >= reader.getNSE() && "Not enough space in buffers"); \ + assert(cSize >= lvlRank * reader.getNSE()); \ + assert(vSize >= reader.getNSE()); \ (void)dimRank; \ (void)cSize; \ (void)vSize; \ @@ -488,10 +455,6 @@ index_type sparseDimSize(void *tensor, index_type d) { return static_cast<SparseTensorStorageBase *>(tensor)->getDimSize(d); } -void endForwardingInsert(void *tensor) { - return static_cast<SparseTensorStorageBase *>(tensor)->endForwardingInsert(); -} - void endLexInsert(void *tensor) { return static_cast<SparseTensorStorageBase *>(tensor)->endLexInsert(); } @@ -500,13 +463,6 @@ void delSparseTensor(void *tensor) { delete static_cast<SparseTensorStorageBase *>(tensor); } -#define IMPL_DELCOO(VNAME, V) \ - void delSparseTensorCOO##VNAME(void *coo) { \ - delete static_cast<SparseTensorCOO<V> *>(coo); \ - } -MLIR_SPARSETENSOR_FOREVERY_V(IMPL_DELCOO) -#undef IMPL_DELCOO - char *getTensorFilename(index_type id) { constexpr size_t bufSize = 80; char var[bufSize]; diff --git a/mlir/lib/Target/LLVMIR/DebugImporter.cpp b/mlir/lib/Target/LLVMIR/DebugImporter.cpp index 6521295..c631617 100644 --- a/mlir/lib/Target/LLVMIR/DebugImporter.cpp +++ b/mlir/lib/Target/LLVMIR/DebugImporter.cpp @@ -99,21 +99,31 @@ DIFileAttr DebugImporter::translateImpl(llvm::DIFile *node) { } DILabelAttr DebugImporter::translateImpl(llvm::DILabel *node) { - return DILabelAttr::get(context, translate(node->getScope()), + // Return nullptr if the scope or type is a cyclic dependency. + DIScopeAttr scope = translate(node->getScope()); + if (node->getScope() && !scope) + return nullptr; + return DILabelAttr::get(context, scope, getStringAttrOrNull(node->getRawName()), translate(node->getFile()), node->getLine()); } DILexicalBlockAttr DebugImporter::translateImpl(llvm::DILexicalBlock *node) { - return DILexicalBlockAttr::get(context, translate(node->getScope()), - translate(node->getFile()), node->getLine(), - node->getColumn()); + // Return nullptr if the scope or type is a cyclic dependency. + DIScopeAttr scope = translate(node->getScope()); + if (node->getScope() && !scope) + return nullptr; + return DILexicalBlockAttr::get(context, scope, translate(node->getFile()), + node->getLine(), node->getColumn()); } DILexicalBlockFileAttr DebugImporter::translateImpl(llvm::DILexicalBlockFile *node) { - return DILexicalBlockFileAttr::get(context, translate(node->getScope()), - translate(node->getFile()), + // Return nullptr if the scope or type is a cyclic dependency. + DIScopeAttr scope = translate(node->getScope()); + if (node->getScope() && !scope) + return nullptr; + return DILexicalBlockFileAttr::get(context, scope, translate(node->getFile()), node->getDiscriminator()); } @@ -135,11 +145,14 @@ DebugImporter::translateImpl(llvm::DIGlobalVariable *node) { } DILocalVariableAttr DebugImporter::translateImpl(llvm::DILocalVariable *node) { - return DILocalVariableAttr::get(context, translate(node->getScope()), - getStringAttrOrNull(node->getRawName()), - translate(node->getFile()), node->getLine(), - node->getArg(), node->getAlignInBits(), - translate(node->getType())); + // Return nullptr if the scope or type is a cyclic dependency. + DIScopeAttr scope = translate(node->getScope()); + if (node->getScope() && !scope) + return nullptr; + return DILocalVariableAttr::get( + context, scope, getStringAttrOrNull(node->getRawName()), + translate(node->getFile()), node->getLine(), node->getArg(), + node->getAlignInBits(), translate(node->getType())); } DIScopeAttr DebugImporter::translateImpl(llvm::DIScope *node) { diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp index 97ccb2b..d63ea12 100644 --- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp @@ -1966,6 +1966,13 @@ ModuleImport::processDebugIntrinsic(llvm::DbgVariableIntrinsic *dbgIntr, // TODO: find a way to support this case. if (isMetadataKillLocation(dbgIntr)) return emitUnsupportedWarning(); + // Drop debug intrinsics if the associated variable information cannot be + // translated due to cyclic debug metadata. + // TODO: Support cyclic debug metadata. + DILocalVariableAttr localVariableAttr = + matchLocalVariableAttr(dbgIntr->getArgOperand(1)); + if (!localVariableAttr) + return emitUnsupportedWarning(); FailureOr<Value> argOperand = convertMetadataValue(dbgIntr->getArgOperand(0)); if (failed(argOperand)) return emitError(loc) << "failed to convert a debug intrinsic operand: " @@ -1991,8 +1998,6 @@ ModuleImport::processDebugIntrinsic(llvm::DbgVariableIntrinsic *dbgIntr, } else { builder.setInsertionPointAfterValue(*argOperand); } - DILocalVariableAttr localVariableAttr = - matchLocalVariableAttr(dbgIntr->getArgOperand(1)); auto locationExprAttr = debugImporter->translateExpression(dbgIntr->getExpression()); Operation *op = diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp index 4989ddc..857b601 100644 --- a/mlir/lib/Transforms/Utils/DialectConversion.cpp +++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp @@ -152,519 +152,25 @@ namespace { /// This class contains a snapshot of the current conversion rewriter state. /// This is useful when saving and undoing a set of rewrites. struct RewriterState { - RewriterState(unsigned numCreatedOps, unsigned numUnresolvedMaterializations, - unsigned numReplacements, unsigned numArgReplacements, - unsigned numRewrites, unsigned numIgnoredOperations) - : numCreatedOps(numCreatedOps), - numUnresolvedMaterializations(numUnresolvedMaterializations), - numReplacements(numReplacements), - numArgReplacements(numArgReplacements), numRewrites(numRewrites), - numIgnoredOperations(numIgnoredOperations) {} - - /// The current number of created operations. - unsigned numCreatedOps; - - /// The current number of unresolved materializations. - unsigned numUnresolvedMaterializations; - - /// The current number of replacements queued. - unsigned numReplacements; - - /// The current number of argument replacements queued. - unsigned numArgReplacements; + RewriterState(unsigned numRewrites, unsigned numIgnoredOperations, + unsigned numErased) + : numRewrites(numRewrites), numIgnoredOperations(numIgnoredOperations), + numErased(numErased) {} /// The current number of rewrites performed. unsigned numRewrites; /// The current number of ignored operations. unsigned numIgnoredOperations; -}; - -//===----------------------------------------------------------------------===// -// OpReplacement - -/// This class represents one requested operation replacement via 'replaceOp' or -/// 'eraseOp`. -struct OpReplacement { - OpReplacement(const TypeConverter *converter = nullptr) - : converter(converter) {} - - /// An optional type converter that can be used to materialize conversions - /// between the new and old values if necessary. - const TypeConverter *converter; -}; - -//===----------------------------------------------------------------------===// -// UnresolvedMaterialization - -/// This class represents an unresolved materialization, i.e. a materialization -/// that was inserted during conversion that needs to be legalized at the end of -/// the conversion process. -class UnresolvedMaterialization { -public: - /// The type of materialization. - enum Kind { - /// This materialization materializes a conversion for an illegal block - /// argument type, to a legal one. - Argument, - - /// This materialization materializes a conversion from an illegal type to a - /// legal one. - Target - }; - - UnresolvedMaterialization(UnrealizedConversionCastOp op = nullptr, - const TypeConverter *converter = nullptr, - Kind kind = Target, Type origOutputType = nullptr) - : op(op), converterAndKind(converter, kind), - origOutputType(origOutputType) {} - - /// Return the temporary conversion operation inserted for this - /// materialization. - UnrealizedConversionCastOp getOp() const { return op; } - /// Return the type converter of this materialization (which may be null). - const TypeConverter *getConverter() const { - return converterAndKind.getPointer(); - } - - /// Return the kind of this materialization. - Kind getKind() const { return converterAndKind.getInt(); } - - /// Set the kind of this materialization. - void setKind(Kind kind) { converterAndKind.setInt(kind); } - - /// Return the original illegal output type of the input values. - Type getOrigOutputType() const { return origOutputType; } - -private: - /// The unresolved materialization operation created during conversion. - UnrealizedConversionCastOp op; - - /// The corresponding type converter to use when resolving this - /// materialization, and the kind of this materialization. - llvm::PointerIntPair<const TypeConverter *, 1, Kind> converterAndKind; - - /// The original output type. This is only used for argument conversions. - Type origOutputType; + /// The current number of erased operations/blocks. + unsigned numErased; }; -} // namespace - -/// Build an unresolved materialization operation given an output type and set -/// of input operands. -static Value buildUnresolvedMaterialization( - UnresolvedMaterialization::Kind kind, Block *insertBlock, - Block::iterator insertPt, Location loc, ValueRange inputs, Type outputType, - Type origOutputType, const TypeConverter *converter, - SmallVectorImpl<UnresolvedMaterialization> &unresolvedMaterializations) { - // Avoid materializing an unnecessary cast. - if (inputs.size() == 1 && inputs.front().getType() == outputType) - return inputs.front(); - - // Create an unresolved materialization. We use a new OpBuilder to avoid - // tracking the materialization like we do for other operations. - OpBuilder builder(insertBlock, insertPt); - auto convertOp = - builder.create<UnrealizedConversionCastOp>(loc, outputType, inputs); - unresolvedMaterializations.emplace_back(convertOp, converter, kind, - origOutputType); - return convertOp.getResult(0); -} -static Value buildUnresolvedArgumentMaterialization( - PatternRewriter &rewriter, Location loc, ValueRange inputs, - Type origOutputType, Type outputType, const TypeConverter *converter, - SmallVectorImpl<UnresolvedMaterialization> &unresolvedMaterializations) { - return buildUnresolvedMaterialization( - UnresolvedMaterialization::Argument, rewriter.getInsertionBlock(), - rewriter.getInsertionPoint(), loc, inputs, outputType, origOutputType, - converter, unresolvedMaterializations); -} -static Value buildUnresolvedTargetMaterialization( - Location loc, Value input, Type outputType, const TypeConverter *converter, - SmallVectorImpl<UnresolvedMaterialization> &unresolvedMaterializations) { - Block *insertBlock = input.getParentBlock(); - Block::iterator insertPt = insertBlock->begin(); - if (OpResult inputRes = dyn_cast<OpResult>(input)) - insertPt = ++inputRes.getOwner()->getIterator(); - - return buildUnresolvedMaterialization( - UnresolvedMaterialization::Target, insertBlock, insertPt, loc, input, - outputType, outputType, converter, unresolvedMaterializations); -} - -//===----------------------------------------------------------------------===// -// ArgConverter -//===----------------------------------------------------------------------===// -namespace { -/// This class provides a simple interface for converting the types of block -/// arguments. This is done by creating a new block that contains the new legal -/// types and extracting the block that contains the old illegal types to allow -/// for undoing pending rewrites in the case of failure. -struct ArgConverter { - ArgConverter( - PatternRewriter &rewriter, - SmallVectorImpl<UnresolvedMaterialization> &unresolvedMaterializations) - : rewriter(rewriter), - unresolvedMaterializations(unresolvedMaterializations) {} - - /// This structure contains the information pertaining to an argument that has - /// been converted. - struct ConvertedArgInfo { - ConvertedArgInfo(unsigned newArgIdx, unsigned newArgSize, - Value castValue = nullptr) - : newArgIdx(newArgIdx), newArgSize(newArgSize), castValue(castValue) {} - - /// The start index of in the new argument list that contains arguments that - /// replace the original. - unsigned newArgIdx; - - /// The number of arguments that replaced the original argument. - unsigned newArgSize; - - /// The cast value that was created to cast from the new arguments to the - /// old. This only used if 'newArgSize' > 1. - Value castValue; - }; - - /// This structure contains information pertaining to a block that has had its - /// signature converted. - struct ConvertedBlockInfo { - ConvertedBlockInfo(Block *origBlock, const TypeConverter *converter) - : origBlock(origBlock), converter(converter) {} - - /// The original block that was requested to have its signature converted. - Block *origBlock; - - /// The conversion information for each of the arguments. The information is - /// std::nullopt if the argument was dropped during conversion. - SmallVector<std::optional<ConvertedArgInfo>, 1> argInfo; - - /// The type converter used to convert the arguments. - const TypeConverter *converter; - }; - - //===--------------------------------------------------------------------===// - // Rewrite Application - //===--------------------------------------------------------------------===// - - /// Erase any rewrites registered for the blocks within the given operation - /// which is about to be removed. This merely drops the rewrites without - /// undoing them. - void notifyOpRemoved(Operation *op); - - /// Cleanup and undo any generated conversions for the arguments of block. - /// This method replaces the new block with the original, reverting the IR to - /// its original state. - void discardRewrites(Block *block); - - /// Fully replace uses of the old arguments with the new. - void applyRewrites(ConversionValueMapping &mapping); - - /// Materialize any necessary conversions for converted arguments that have - /// live users, using the provided `findLiveUser` to search for a user that - /// survives the conversion process. - LogicalResult - materializeLiveConversions(ConversionValueMapping &mapping, - OpBuilder &builder, - function_ref<Operation *(Value)> findLiveUser); - - //===--------------------------------------------------------------------===// - // Conversion - //===--------------------------------------------------------------------===// - - /// Attempt to convert the signature of the given block, if successful a new - /// block is returned containing the new arguments. Returns `block` if it did - /// not require conversion. - FailureOr<Block *> - convertSignature(Block *block, const TypeConverter *converter, - ConversionValueMapping &mapping, - SmallVectorImpl<BlockArgument> &argReplacements); - - /// Apply the given signature conversion on the given block. The new block - /// containing the updated signature is returned. If no conversions were - /// necessary, e.g. if the block has no arguments, `block` is returned. - /// `converter` is used to generate any necessary cast operations that - /// translate between the origin argument types and those specified in the - /// signature conversion. - Block *applySignatureConversion( - Block *block, const TypeConverter *converter, - TypeConverter::SignatureConversion &signatureConversion, - ConversionValueMapping &mapping, - SmallVectorImpl<BlockArgument> &argReplacements); - - /// A collection of blocks that have had their arguments converted. This is a - /// map from the new replacement block, back to the original block. - llvm::MapVector<Block *, ConvertedBlockInfo> conversionInfo; - - /// The pattern rewriter to use when materializing conversions. - PatternRewriter &rewriter; - - /// An ordered set of unresolved materializations during conversion. - SmallVectorImpl<UnresolvedMaterialization> &unresolvedMaterializations; -}; -} // namespace - -//===----------------------------------------------------------------------===// -// Rewrite Application - -void ArgConverter::notifyOpRemoved(Operation *op) { - if (conversionInfo.empty()) - return; - - for (Region ®ion : op->getRegions()) { - for (Block &block : region) { - // Drop any rewrites from within. - for (Operation &nestedOp : block) - if (nestedOp.getNumRegions()) - notifyOpRemoved(&nestedOp); - - // Check if this block was converted. - auto *it = conversionInfo.find(&block); - if (it == conversionInfo.end()) - continue; - - // Drop all uses of the original arguments and delete the original block. - Block *origBlock = it->second.origBlock; - for (BlockArgument arg : origBlock->getArguments()) - arg.dropAllUses(); - conversionInfo.erase(it); - } - } -} - -void ArgConverter::discardRewrites(Block *block) { - auto *it = conversionInfo.find(block); - if (it == conversionInfo.end()) - return; - Block *origBlock = it->second.origBlock; - - // Drop all uses of the new block arguments and replace uses of the new block. - for (int i = block->getNumArguments() - 1; i >= 0; --i) - block->getArgument(i).dropAllUses(); - block->replaceAllUsesWith(origBlock); - - // Move the operations back the original block, move the original block back - // into its original location and the delete the new block. - origBlock->getOperations().splice(origBlock->end(), block->getOperations()); - block->getParent()->getBlocks().insert(Region::iterator(block), origBlock); - block->erase(); - - conversionInfo.erase(it); -} - -void ArgConverter::applyRewrites(ConversionValueMapping &mapping) { - for (auto &info : conversionInfo) { - ConvertedBlockInfo &blockInfo = info.second; - Block *origBlock = blockInfo.origBlock; - - // Process the remapping for each of the original arguments. - for (unsigned i = 0, e = origBlock->getNumArguments(); i != e; ++i) { - std::optional<ConvertedArgInfo> &argInfo = blockInfo.argInfo[i]; - BlockArgument origArg = origBlock->getArgument(i); - - // Handle the case of a 1->0 value mapping. - if (!argInfo) { - if (Value newArg = mapping.lookupOrNull(origArg, origArg.getType())) - origArg.replaceAllUsesWith(newArg); - continue; - } - - // Otherwise this is a 1->1+ value mapping. - Value castValue = argInfo->castValue; - assert(argInfo->newArgSize >= 1 && castValue && "expected 1->1+ mapping"); - - // If the argument is still used, replace it with the generated cast. - if (!origArg.use_empty()) { - origArg.replaceAllUsesWith( - mapping.lookupOrDefault(castValue, origArg.getType())); - } - } - - delete origBlock; - blockInfo.origBlock = nullptr; - } -} - -LogicalResult ArgConverter::materializeLiveConversions( - ConversionValueMapping &mapping, OpBuilder &builder, - function_ref<Operation *(Value)> findLiveUser) { - for (auto &info : conversionInfo) { - Block *newBlock = info.first; - ConvertedBlockInfo &blockInfo = info.second; - Block *origBlock = blockInfo.origBlock; - - // Process the remapping for each of the original arguments. - for (unsigned i = 0, e = origBlock->getNumArguments(); i != e; ++i) { - // If the type of this argument changed and the argument is still live, we - // need to materialize a conversion. - BlockArgument origArg = origBlock->getArgument(i); - if (mapping.lookupOrNull(origArg, origArg.getType())) - continue; - Operation *liveUser = findLiveUser(origArg); - if (!liveUser) - continue; - - Value replacementValue = mapping.lookupOrDefault(origArg); - bool isDroppedArg = replacementValue == origArg; - if (isDroppedArg) - rewriter.setInsertionPointToStart(newBlock); - else - rewriter.setInsertionPointAfterValue(replacementValue); - Value newArg; - if (blockInfo.converter) { - newArg = blockInfo.converter->materializeSourceConversion( - rewriter, origArg.getLoc(), origArg.getType(), - isDroppedArg ? ValueRange() : ValueRange(replacementValue)); - assert((!newArg || newArg.getType() == origArg.getType()) && - "materialization hook did not provide a value of the expected " - "type"); - } - if (!newArg) { - InFlightDiagnostic diag = - emitError(origArg.getLoc()) - << "failed to materialize conversion for block argument #" << i - << " that remained live after conversion, type was " - << origArg.getType(); - if (!isDroppedArg) - diag << ", with target type " << replacementValue.getType(); - diag.attachNote(liveUser->getLoc()) - << "see existing live user here: " << *liveUser; - return failure(); - } - mapping.map(origArg, newArg); - } - } - return success(); -} - -//===----------------------------------------------------------------------===// -// Conversion - -FailureOr<Block *> ArgConverter::convertSignature( - Block *block, const TypeConverter *converter, - ConversionValueMapping &mapping, - SmallVectorImpl<BlockArgument> &argReplacements) { - // Check if the block was already converted. - // * If the block is mapped in `conversionInfo`, it is a converted block. - // * If the block is detached, conservatively assume that it is going to be - // deleted; it is likely the old block (before it was converted). - if (conversionInfo.count(block) || !block->getParent()) - return block; - // If a converter wasn't provided, and the block wasn't already converted, - // there is nothing we can do. - if (!converter) - return failure(); - - // Try to convert the signature for the block with the provided converter. - if (auto conversion = converter->convertBlockSignature(block)) - return applySignatureConversion(block, converter, *conversion, mapping, - argReplacements); - return failure(); -} - -Block *ArgConverter::applySignatureConversion( - Block *block, const TypeConverter *converter, - TypeConverter::SignatureConversion &signatureConversion, - ConversionValueMapping &mapping, - SmallVectorImpl<BlockArgument> &argReplacements) { - // If no arguments are being changed or added, there is nothing to do. - unsigned origArgCount = block->getNumArguments(); - auto convertedTypes = signatureConversion.getConvertedTypes(); - if (origArgCount == 0 && convertedTypes.empty()) - return block; - - // Split the block at the beginning to get a new block to use for the updated - // signature. - Block *newBlock = block->splitBlock(block->begin()); - block->replaceAllUsesWith(newBlock); - // Unlink the block, but do not erase it yet, so that the change can be rolled - // back. - block->getParent()->getBlocks().remove(block); - - // Map all new arguments to the location of the argument they originate from. - SmallVector<Location> newLocs(convertedTypes.size(), - rewriter.getUnknownLoc()); - for (unsigned i = 0; i < origArgCount; ++i) { - auto inputMap = signatureConversion.getInputMapping(i); - if (!inputMap || inputMap->replacementValue) - continue; - Location origLoc = block->getArgument(i).getLoc(); - for (unsigned j = 0; j < inputMap->size; ++j) - newLocs[inputMap->inputNo + j] = origLoc; - } - - SmallVector<Value, 4> newArgRange( - newBlock->addArguments(convertedTypes, newLocs)); - ArrayRef<Value> newArgs(newArgRange); - - // Remap each of the original arguments as determined by the signature - // conversion. - ConvertedBlockInfo info(block, converter); - info.argInfo.resize(origArgCount); - - OpBuilder::InsertionGuard guard(rewriter); - rewriter.setInsertionPointToStart(newBlock); - for (unsigned i = 0; i != origArgCount; ++i) { - auto inputMap = signatureConversion.getInputMapping(i); - if (!inputMap) - continue; - BlockArgument origArg = block->getArgument(i); - - // If inputMap->replacementValue is not nullptr, then the argument is - // dropped and a replacement value is provided to be the remappedValue. - if (inputMap->replacementValue) { - assert(inputMap->size == 0 && - "invalid to provide a replacement value when the argument isn't " - "dropped"); - mapping.map(origArg, inputMap->replacementValue); - argReplacements.push_back(origArg); - continue; - } - - // Otherwise, this is a 1->1+ mapping. - auto replArgs = newArgs.slice(inputMap->inputNo, inputMap->size); - Value newArg; - - // If this is a 1->1 mapping and the types of new and replacement arguments - // match (i.e. it's an identity map), then the argument is mapped to its - // original type. - // FIXME: We simply pass through the replacement argument if there wasn't a - // converter, which isn't great as it allows implicit type conversions to - // appear. We should properly restructure this code to handle cases where a - // converter isn't provided and also to properly handle the case where an - // argument materialization is actually a temporary source materialization - // (e.g. in the case of 1->N). - if (replArgs.size() == 1 && - (!converter || replArgs[0].getType() == origArg.getType())) { - newArg = replArgs.front(); - } else { - Type origOutputType = origArg.getType(); - - // Legalize the argument output type. - Type outputType = origOutputType; - if (Type legalOutputType = converter->convertType(outputType)) - outputType = legalOutputType; - - newArg = buildUnresolvedArgumentMaterialization( - rewriter, origArg.getLoc(), replArgs, origOutputType, outputType, - converter, unresolvedMaterializations); - } - - mapping.map(origArg, newArg); - argReplacements.push_back(origArg); - info.argInfo[i] = - ConvertedArgInfo(inputMap->inputNo, inputMap->size, newArg); - } - - conversionInfo.insert({newBlock, std::move(info)}); - return newBlock; -} //===----------------------------------------------------------------------===// // IR rewrites //===----------------------------------------------------------------------===// -namespace { /// An IR rewrite that can be committed (upon success) or rolled back (upon /// failure). /// @@ -685,19 +191,29 @@ public: MoveBlock, SplitBlock, BlockTypeConversion, + ReplaceBlockArg, // Operation rewrites MoveOperation, - ModifyOperation + ModifyOperation, + ReplaceOperation, + CreateOperation, + UnresolvedMaterialization }; virtual ~IRRewrite() = default; - /// Roll back the rewrite. + /// Roll back the rewrite. Operations may be erased during rollback. virtual void rollback() = 0; - /// Commit the rewrite. + /// Commit the rewrite. Operations may be unlinked from their blocks during + /// the commit phase, but they must not be erased yet. This is because + /// internal dialect conversion state (such as `mapping`) may still be using + /// them. Operations must be erased during cleanup. virtual void commit() {} + /// Cleanup operations. Cleanup is called after commit. + virtual void cleanup() {} + Kind getKind() const { return kind; } static bool classof(const IRRewrite *rewrite) { return true; } @@ -706,6 +222,14 @@ protected: IRRewrite(Kind kind, ConversionPatternRewriterImpl &rewriterImpl) : kind(kind), rewriterImpl(rewriterImpl) {} + /// Erase the given op (unless it was already erased). + void eraseOp(Operation *op); + + /// Erase the given block (unless it was already erased). + void eraseBlock(Block *block); + + const ConversionConfig &getConfig() const; + const Kind kind; ConversionPatternRewriterImpl &rewriterImpl; }; @@ -718,7 +242,7 @@ public: static bool classof(const IRRewrite *rewrite) { return rewrite->getKind() >= Kind::CreateBlock && - rewrite->getKind() <= Kind::BlockTypeConversion; + rewrite->getKind() <= Kind::ReplaceBlockArg; } protected: @@ -748,8 +272,10 @@ public: auto &blockOps = block->getOperations(); while (!blockOps.empty()) blockOps.remove(blockOps.begin()); - block->dropAllDefinedValueUses(); - block->erase(); + if (block->getParent()) + eraseBlock(block); + else + delete block; } }; @@ -787,6 +313,8 @@ public: void commit() override { // Erase the block. assert(block && "expected block"); + assert(block->empty() && "expected empty block"); + block->dropAllDefinedValueUses(); delete block; block = nullptr; } @@ -885,8 +413,7 @@ public: // Merge back the block that was split out. originalBlock->getOperations().splice(originalBlock->end(), block->getOperations()); - block->dropAllDefinedValueUses(); - block->erase(); + eraseBlock(block); } private: @@ -894,20 +421,80 @@ private: Block *originalBlock; }; +/// This structure contains the information pertaining to an argument that has +/// been converted. +struct ConvertedArgInfo { + ConvertedArgInfo(unsigned newArgIdx, unsigned newArgSize, + Value castValue = nullptr) + : newArgIdx(newArgIdx), newArgSize(newArgSize), castValue(castValue) {} + + /// The start index of in the new argument list that contains arguments that + /// replace the original. + unsigned newArgIdx; + + /// The number of arguments that replaced the original argument. + unsigned newArgSize; + + /// The cast value that was created to cast from the new arguments to the + /// old. This only used if 'newArgSize' > 1. + Value castValue; +}; + /// Block type conversion. This rewrite is partially reflected in the IR. class BlockTypeConversionRewrite : public BlockRewrite { public: - BlockTypeConversionRewrite(ConversionPatternRewriterImpl &rewriterImpl, - Block *block) - : BlockRewrite(Kind::BlockTypeConversion, rewriterImpl, block) {} + BlockTypeConversionRewrite( + ConversionPatternRewriterImpl &rewriterImpl, Block *block, + Block *origBlock, SmallVector<std::optional<ConvertedArgInfo>, 1> argInfo, + const TypeConverter *converter) + : BlockRewrite(Kind::BlockTypeConversion, rewriterImpl, block), + origBlock(origBlock), argInfo(argInfo), converter(converter) {} static bool classof(const IRRewrite *rewrite) { return rewrite->getKind() == Kind::BlockTypeConversion; } - // TODO: Block type conversions are currently committed in - // `ArgConverter::applyRewrites`. This should be done in the "commit" method. + /// Materialize any necessary conversions for converted arguments that have + /// live users, using the provided `findLiveUser` to search for a user that + /// survives the conversion process. + LogicalResult + materializeLiveConversions(function_ref<Operation *(Value)> findLiveUser); + + void commit() override; + void rollback() override; + +private: + /// The original block that was requested to have its signature converted. + Block *origBlock; + + /// The conversion information for each of the arguments. The information is + /// std::nullopt if the argument was dropped during conversion. + SmallVector<std::optional<ConvertedArgInfo>, 1> argInfo; + + /// The type converter used to convert the arguments. + const TypeConverter *converter; +}; + +/// Replacing a block argument. This rewrite is not immediately reflected in the +/// IR. An internal IR mapping is updated, but the actual replacement is delayed +/// until the rewrite is committed. +class ReplaceBlockArgRewrite : public BlockRewrite { +public: + ReplaceBlockArgRewrite(ConversionPatternRewriterImpl &rewriterImpl, + Block *block, BlockArgument arg) + : BlockRewrite(Kind::ReplaceBlockArg, rewriterImpl, block), arg(arg) {} + + static bool classof(const IRRewrite *rewrite) { + return rewrite->getKind() == Kind::ReplaceBlockArg; + } + + void commit() override; + + void rollback() override; + +private: + BlockArgument arg; }; /// An operation rewrite. @@ -918,7 +505,7 @@ public: static bool classof(const IRRewrite *rewrite) { return rewrite->getKind() >= Kind::MoveOperation && - rewrite->getKind() <= Kind::ModifyOperation; + rewrite->getKind() <= Kind::UnresolvedMaterialization; } protected: @@ -953,8 +540,8 @@ private: // The block in which this operation was previously contained. Block *block; - // The original successor of this operation before it was moved. "nullptr" if - // this operation was the only operation in the region. + // The original successor of this operation before it was moved. "nullptr" + // if this operation was the only operation in the region. Operation *insertBeforeOp; }; @@ -1019,6 +606,118 @@ private: SmallVector<Block *, 2> successors; void *propertiesStorage = nullptr; }; + +/// Replacing an operation. Erasing an operation is treated as a special case +/// with "null" replacements. This rewrite is not immediately reflected in the +/// IR. An internal IR mapping is updated, but values are not replaced and the +/// original op is not erased until the rewrite is committed. +class ReplaceOperationRewrite : public OperationRewrite { +public: + ReplaceOperationRewrite(ConversionPatternRewriterImpl &rewriterImpl, + Operation *op, const TypeConverter *converter, + bool changedResults) + : OperationRewrite(Kind::ReplaceOperation, rewriterImpl, op), + converter(converter), changedResults(changedResults) {} + + static bool classof(const IRRewrite *rewrite) { + return rewrite->getKind() == Kind::ReplaceOperation; + } + + void commit() override; + + void rollback() override; + + void cleanup() override; + + const TypeConverter *getConverter() const { return converter; } + + bool hasChangedResults() const { return changedResults; } + +private: + /// An optional type converter that can be used to materialize conversions + /// between the new and old values if necessary. + const TypeConverter *converter; + + /// A boolean flag that indicates whether result types have changed or not. + bool changedResults; +}; + +class CreateOperationRewrite : public OperationRewrite { +public: + CreateOperationRewrite(ConversionPatternRewriterImpl &rewriterImpl, + Operation *op) + : OperationRewrite(Kind::CreateOperation, rewriterImpl, op) {} + + static bool classof(const IRRewrite *rewrite) { + return rewrite->getKind() == Kind::CreateOperation; + } + + void rollback() override; +}; + +/// The type of materialization. +enum MaterializationKind { + /// This materialization materializes a conversion for an illegal block + /// argument type, to a legal one. + Argument, + + /// This materialization materializes a conversion from an illegal type to a + /// legal one. + Target +}; + +/// An unresolved materialization, i.e., a "builtin.unrealized_conversion_cast" +/// op. Unresolved materializations are erased at the end of the dialect +/// conversion. +class UnresolvedMaterializationRewrite : public OperationRewrite { +public: + UnresolvedMaterializationRewrite( + ConversionPatternRewriterImpl &rewriterImpl, + UnrealizedConversionCastOp op, const TypeConverter *converter = nullptr, + MaterializationKind kind = MaterializationKind::Target, + Type origOutputType = nullptr) + : OperationRewrite(Kind::UnresolvedMaterialization, rewriterImpl, op), + converterAndKind(converter, kind), origOutputType(origOutputType) {} + + static bool classof(const IRRewrite *rewrite) { + return rewrite->getKind() == Kind::UnresolvedMaterialization; + } + + UnrealizedConversionCastOp getOperation() const { + return cast<UnrealizedConversionCastOp>(op); + } + + void rollback() override; + + void cleanup() override; + + /// Return the type converter of this materialization (which may be null). + const TypeConverter *getConverter() const { + return converterAndKind.getPointer(); + } + + /// Return the kind of this materialization. + MaterializationKind getMaterializationKind() const { + return converterAndKind.getInt(); + } + + /// Set the kind of this materialization. + void setMaterializationKind(MaterializationKind kind) { + converterAndKind.setInt(kind); + } + + /// Return the original illegal output type of the input values. + Type getOrigOutputType() const { return origOutputType; } + +private: + /// The corresponding type converter to use when resolving this + /// materialization, and the kind of this materialization. + llvm::PointerIntPair<const TypeConverter *, 1, MaterializationKind> + converterAndKind; + + /// The original output type. This is only used for argument conversions. + Type origOutputType; +}; } // namespace /// Return "true" if there is an operation rewrite that matches the specified @@ -1031,23 +730,35 @@ static bool hasRewrite(R &&rewrites, Operation *op) { }); } +/// Find the single rewrite object of the specified type and block among the +/// given rewrites. In debug mode, asserts that there is mo more than one such +/// object. Return "nullptr" if no object was found. +template <typename RewriteTy, typename R> +static RewriteTy *findSingleRewrite(R &&rewrites, Block *block) { + RewriteTy *result = nullptr; + for (auto &rewrite : rewrites) { + auto *rewriteTy = dyn_cast<RewriteTy>(rewrite.get()); + if (rewriteTy && rewriteTy->getBlock() == block) { +#ifndef NDEBUG + assert(!result && "expected single matching rewrite"); + result = rewriteTy; +#else + return rewriteTy; +#endif // NDEBUG + } + } + return result; +} + //===----------------------------------------------------------------------===// // ConversionPatternRewriterImpl //===----------------------------------------------------------------------===// namespace mlir { namespace detail { struct ConversionPatternRewriterImpl : public RewriterBase::Listener { - explicit ConversionPatternRewriterImpl(PatternRewriter &rewriter) - : argConverter(rewriter, unresolvedMaterializations), - notifyCallback(nullptr) {} - - /// Cleanup and destroy any generated rewrite operations. This method is - /// invoked when the conversion process fails. - void discardRewrites(); - - /// Apply all requested operation rewrites. This method is invoked when the - /// conversion process succeeds. - void applyRewrites(); + explicit ConversionPatternRewriterImpl(MLIRContext *ctx, + const ConversionConfig &config) + : eraseRewriter(ctx), config(config) {} //===--------------------------------------------------------------------===// // State Management @@ -1056,6 +767,10 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener { /// Return the current state of the rewriter. RewriterState getCurrentState(); + /// Apply all requested operation rewrites. This method is invoked when the + /// conversion process succeeds. + void applyRewrites(); + /// Reset the state of the rewriter to a previously saved point. void resetState(RewriterState state); @@ -1092,11 +807,18 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener { // Type Conversion //===--------------------------------------------------------------------===// - /// Convert the signature of the given block. + /// Attempt to convert the signature of the given block, if successful a new + /// block is returned containing the new arguments. Returns `block` if it did + /// not require conversion. FailureOr<Block *> convertBlockSignature( Block *block, const TypeConverter *converter, TypeConverter::SignatureConversion *conversion = nullptr); + /// Convert the types of non-entry block arguments within the given region. + LogicalResult convertNonEntryRegionTypes( + Region *region, const TypeConverter &converter, + ArrayRef<TypeConverter::SignatureConversion> blockConversions = {}); + /// Apply a signature conversion on the given region, using `converter` for /// materializations if not null. Block * @@ -1109,10 +831,37 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener { convertRegionTypes(Region *region, const TypeConverter &converter, TypeConverter::SignatureConversion *entryConversion); - /// Convert the types of non-entry block arguments within the given region. - LogicalResult convertNonEntryRegionTypes( - Region *region, const TypeConverter &converter, - ArrayRef<TypeConverter::SignatureConversion> blockConversions = {}); + /// Apply the given signature conversion on the given block. The new block + /// containing the updated signature is returned. If no conversions were + /// necessary, e.g. if the block has no arguments, `block` is returned. + /// `converter` is used to generate any necessary cast operations that + /// translate between the origin argument types and those specified in the + /// signature conversion. + Block *applySignatureConversion( + Block *block, const TypeConverter *converter, + TypeConverter::SignatureConversion &signatureConversion); + + //===--------------------------------------------------------------------===// + // Materializations + //===--------------------------------------------------------------------===// + /// Build an unresolved materialization operation given an output type and set + /// of input operands. + Value buildUnresolvedMaterialization(MaterializationKind kind, + Block *insertBlock, + Block::iterator insertPt, Location loc, + ValueRange inputs, Type outputType, + Type origOutputType, + const TypeConverter *converter); + + Value buildUnresolvedArgumentMaterialization(Block *block, Location loc, + ValueRange inputs, + Type origOutputType, + Type outputType, + const TypeConverter *converter); + + Value buildUnresolvedTargetMaterialization(Location loc, Value input, + Type outputType, + const TypeConverter *converter); //===--------------------------------------------------------------------===// // Rewriter Notification Hooks @@ -1145,28 +894,51 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener { function_ref<void(Diagnostic &)> reasonCallback) override; //===--------------------------------------------------------------------===// - // State + // IR Erasure //===--------------------------------------------------------------------===// - // Mapping between replaced values that differ in type. This happens when - // replacing a value with one of a different type. - ConversionValueMapping mapping; + /// A rewriter that keeps track of erased ops and blocks. It ensures that no + /// operation or block is erased multiple times. This rewriter assumes that + /// no new IR is created between calls to `eraseOp`/`eraseBlock`. + struct SingleEraseRewriter : public RewriterBase, RewriterBase::Listener { + public: + SingleEraseRewriter(MLIRContext *context) + : RewriterBase(context, /*listener=*/this) {} + + /// Erase the given op (unless it was already erased). + void eraseOp(Operation *op) override { + if (erased.contains(op)) + return; + op->dropAllUses(); + RewriterBase::eraseOp(op); + } - /// Utility used to convert block arguments. - ArgConverter argConverter; + /// Erase the given block (unless it was already erased). + void eraseBlock(Block *block) override { + if (erased.contains(block)) + return; + assert(block->empty() && "expected empty block"); + block->dropAllDefinedValueUses(); + RewriterBase::eraseBlock(block); + } - /// Ordered vector of all of the newly created operations during conversion. - SmallVector<Operation *> createdOps; + void notifyOperationErased(Operation *op) override { erased.insert(op); } + void notifyBlockErased(Block *block) override { erased.insert(block); } - /// Ordered vector of all unresolved type conversion materializations during - /// conversion. - SmallVector<UnresolvedMaterialization> unresolvedMaterializations; + /// Pointers to all erased operations and blocks. + SetVector<void *> erased; + }; + + //===--------------------------------------------------------------------===// + // State + //===--------------------------------------------------------------------===// - /// Ordered map of requested operation replacements. - llvm::MapVector<Operation *, OpReplacement> replacements; + /// This rewriter must be used for erasing ops/blocks. + SingleEraseRewriter eraseRewriter; - /// Ordered vector of any requested block argument replacements. - SmallVector<BlockArgument, 4> argReplacements; + // Mapping between replaced values that differ in type. This happens when + // replacing a value with one of a different type. + ConversionValueMapping mapping; /// Ordered list of block operations (creations, splits, motions). SmallVector<std::unique_ptr<IRRewrite>> rewrites; @@ -1182,11 +954,6 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener { /// operation was ignored. SetVector<Operation *> ignoredOps; - /// A vector of indices into `replacements` of operations that were replaced - /// with values with different result types than the original operation, e.g. - /// 1->N conversion of some kind. - SmallVector<unsigned, 4> operationsWithChangedResults; - /// The current type converter, or nullptr if no type converter is currently /// active. const TypeConverter *currentTypeConverter = nullptr; @@ -1195,8 +962,8 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener { /// converting the arguments of blocks within that region. DenseMap<Region *, const TypeConverter *> regionToConverter; - /// This allows the user to collect the match failure message. - function_ref<void(Diagnostic &)> notifyCallback; + /// Dialect conversion configuration. + const ConversionConfig &config; #ifndef NDEBUG /// A set of operations that have pending updates. This tracking isn't @@ -1211,154 +978,193 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener { } // namespace detail } // namespace mlir -void BlockTypeConversionRewrite::rollback() { - // Undo the type conversion. - rewriterImpl.argConverter.discardRewrites(block); -} - -/// Detach any operations nested in the given operation from their parent -/// blocks, and erase the given operation. This can be used when the nested -/// operations are scheduled for erasure themselves, so deleting the regions of -/// the given operation together with their content would result in double-free. -/// This happens, for example, when rolling back op creation in the reverse -/// order and if the nested ops were created before the parent op. This function -/// does not need to collect nested ops recursively because it is expected to -/// also be called for each nested op when it is about to be deleted. -static void detachNestedAndErase(Operation *op) { - for (Region ®ion : op->getRegions()) { - for (Block &block : region.getBlocks()) { - while (!block.getOperations().empty()) - block.getOperations().remove(block.getOperations().begin()); - block.dropAllDefinedValueUses(); +void IRRewrite::eraseOp(Operation *op) { + rewriterImpl.eraseRewriter.eraseOp(op); +} + +void IRRewrite::eraseBlock(Block *block) { + rewriterImpl.eraseRewriter.eraseBlock(block); +} + +const ConversionConfig &IRRewrite::getConfig() const { + return rewriterImpl.config; +} + +void BlockTypeConversionRewrite::commit() { + // Process the remapping for each of the original arguments. + for (auto [origArg, info] : + llvm::zip_equal(origBlock->getArguments(), argInfo)) { + // Handle the case of a 1->0 value mapping. + if (!info) { + if (Value newArg = + rewriterImpl.mapping.lookupOrNull(origArg, origArg.getType())) + origArg.replaceAllUsesWith(newArg); + continue; + } + + // Otherwise this is a 1->1+ value mapping. + Value castValue = info->castValue; + assert(info->newArgSize >= 1 && castValue && "expected 1->1+ mapping"); + + // If the argument is still used, replace it with the generated cast. + if (!origArg.use_empty()) { + origArg.replaceAllUsesWith( + rewriterImpl.mapping.lookupOrDefault(castValue, origArg.getType())); } } - op->dropAllUses(); - op->erase(); + + assert(origBlock->empty() && "expected empty block"); + origBlock->dropAllDefinedValueUses(); + delete origBlock; + origBlock = nullptr; } -void ConversionPatternRewriterImpl::discardRewrites() { - undoRewrites(); +void BlockTypeConversionRewrite::rollback() { + // Drop all uses of the new block arguments and replace uses of the new block. + for (int i = block->getNumArguments() - 1; i >= 0; --i) + block->getArgument(i).dropAllUses(); + block->replaceAllUsesWith(origBlock); - // Remove any newly created ops. - for (UnresolvedMaterialization &materialization : unresolvedMaterializations) - detachNestedAndErase(materialization.getOp()); - for (auto *op : llvm::reverse(createdOps)) - detachNestedAndErase(op); + // Move the operations back the original block, move the original block back + // into its original location and the delete the new block. + origBlock->getOperations().splice(origBlock->end(), block->getOperations()); + block->getParent()->getBlocks().insert(Region::iterator(block), origBlock); + eraseBlock(block); } -void ConversionPatternRewriterImpl::applyRewrites() { - // Apply all of the rewrites replacements requested during conversion. - for (auto &repl : replacements) { - for (OpResult result : repl.first->getResults()) - if (Value newValue = mapping.lookupOrNull(result, result.getType())) - result.replaceAllUsesWith(newValue); - - // If this operation defines any regions, drop any pending argument - // rewrites. - if (repl.first->getNumRegions()) - argConverter.notifyOpRemoved(repl.first); - } - - // Apply all of the requested argument replacements. - for (BlockArgument arg : argReplacements) { - Value repl = mapping.lookupOrNull(arg, arg.getType()); - if (!repl) +LogicalResult BlockTypeConversionRewrite::materializeLiveConversions( + function_ref<Operation *(Value)> findLiveUser) { + // Process the remapping for each of the original arguments. + for (auto it : llvm::enumerate(origBlock->getArguments())) { + BlockArgument origArg = it.value(); + // Note: `block` may be detached, so OpBuilder::atBlockBegin cannot be used. + OpBuilder builder(it.value().getContext(), /*listener=*/&rewriterImpl); + builder.setInsertionPointToStart(block); + + // If the type of this argument changed and the argument is still live, we + // need to materialize a conversion. + if (rewriterImpl.mapping.lookupOrNull(origArg, origArg.getType())) continue; - - if (isa<BlockArgument>(repl)) { - arg.replaceAllUsesWith(repl); + Operation *liveUser = findLiveUser(origArg); + if (!liveUser) continue; + + Value replacementValue = rewriterImpl.mapping.lookupOrDefault(origArg); + bool isDroppedArg = replacementValue == origArg; + if (!isDroppedArg) + builder.setInsertionPointAfterValue(replacementValue); + Value newArg; + if (converter) { + newArg = converter->materializeSourceConversion( + builder, origArg.getLoc(), origArg.getType(), + isDroppedArg ? ValueRange() : ValueRange(replacementValue)); + assert((!newArg || newArg.getType() == origArg.getType()) && + "materialization hook did not provide a value of the expected " + "type"); } + if (!newArg) { + InFlightDiagnostic diag = + emitError(origArg.getLoc()) + << "failed to materialize conversion for block argument #" + << it.index() << " that remained live after conversion, type was " + << origArg.getType(); + if (!isDroppedArg) + diag << ", with target type " << replacementValue.getType(); + diag.attachNote(liveUser->getLoc()) + << "see existing live user here: " << *liveUser; + return failure(); + } + rewriterImpl.mapping.map(origArg, newArg); + } + return success(); +} - // If the replacement value is an operation, we check to make sure that we - // don't replace uses that are within the parent operation of the - // replacement value. - Operation *replOp = cast<OpResult>(repl).getOwner(); - Block *replBlock = replOp->getBlock(); - arg.replaceUsesWithIf(repl, [&](OpOperand &operand) { - Operation *user = operand.getOwner(); - return user->getBlock() != replBlock || replOp->isBeforeInBlock(user); - }); +void ReplaceBlockArgRewrite::commit() { + Value repl = rewriterImpl.mapping.lookupOrNull(arg, arg.getType()); + if (!repl) + return; + + if (isa<BlockArgument>(repl)) { + arg.replaceAllUsesWith(repl); + return; } - // Drop all of the unresolved materialization operations created during - // conversion. - for (auto &mat : unresolvedMaterializations) { - mat.getOp()->dropAllUses(); - mat.getOp()->erase(); + // If the replacement value is an operation, we check to make sure that we + // don't replace uses that are within the parent operation of the + // replacement value. + Operation *replOp = cast<OpResult>(repl).getOwner(); + Block *replBlock = replOp->getBlock(); + arg.replaceUsesWithIf(repl, [&](OpOperand &operand) { + Operation *user = operand.getOwner(); + return user->getBlock() != replBlock || replOp->isBeforeInBlock(user); + }); +} + +void ReplaceBlockArgRewrite::rollback() { rewriterImpl.mapping.erase(arg); } + +void ReplaceOperationRewrite::commit() { + for (OpResult result : op->getResults()) + if (Value newValue = + rewriterImpl.mapping.lookupOrNull(result, result.getType())) + result.replaceAllUsesWith(newValue); + if (getConfig().unlegalizedOps) + getConfig().unlegalizedOps->erase(op); + // Do not erase the operation yet. It may still be referenced in `mapping`. + op->getBlock()->getOperations().remove(op); +} + +void ReplaceOperationRewrite::rollback() { + for (auto result : op->getResults()) + rewriterImpl.mapping.erase(result); +} + +void ReplaceOperationRewrite::cleanup() { eraseOp(op); } + +void CreateOperationRewrite::rollback() { + for (Region ®ion : op->getRegions()) { + while (!region.getBlocks().empty()) + region.getBlocks().remove(region.getBlocks().begin()); } + op->dropAllUses(); + eraseOp(op); +} - // In a second pass, erase all of the replaced operations in reverse. This - // allows processing nested operations before their parent region is - // destroyed. Because we process in reverse order, producers may be deleted - // before their users (a pattern deleting a producer and then the consumer) - // so we first drop all uses explicitly. - for (auto &repl : llvm::reverse(replacements)) { - repl.first->dropAllUses(); - repl.first->erase(); +void UnresolvedMaterializationRewrite::rollback() { + if (getMaterializationKind() == MaterializationKind::Target) { + for (Value input : op->getOperands()) + rewriterImpl.mapping.erase(input); } + eraseOp(op); +} - argConverter.applyRewrites(mapping); +void UnresolvedMaterializationRewrite::cleanup() { eraseOp(op); } +void ConversionPatternRewriterImpl::applyRewrites() { // Commit all rewrites. for (auto &rewrite : rewrites) rewrite->commit(); + for (auto &rewrite : rewrites) + rewrite->cleanup(); } //===----------------------------------------------------------------------===// // State Management RewriterState ConversionPatternRewriterImpl::getCurrentState() { - return RewriterState(createdOps.size(), unresolvedMaterializations.size(), - replacements.size(), argReplacements.size(), - rewrites.size(), ignoredOps.size()); + return RewriterState(rewrites.size(), ignoredOps.size(), + eraseRewriter.erased.size()); } void ConversionPatternRewriterImpl::resetState(RewriterState state) { - // Reset any replaced arguments. - for (BlockArgument replacedArg : - llvm::drop_begin(argReplacements, state.numArgReplacements)) - mapping.erase(replacedArg); - argReplacements.resize(state.numArgReplacements); - // Undo any rewrites. undoRewrites(state.numRewrites); - // Reset any replaced operations and undo any saved mappings. - for (auto &repl : llvm::drop_begin(replacements, state.numReplacements)) - for (auto result : repl.first->getResults()) - mapping.erase(result); - while (replacements.size() != state.numReplacements) - replacements.pop_back(); - - // Pop all of the newly inserted materializations. - while (unresolvedMaterializations.size() != - state.numUnresolvedMaterializations) { - UnresolvedMaterialization mat = unresolvedMaterializations.pop_back_val(); - UnrealizedConversionCastOp op = mat.getOp(); - - // If this was a target materialization, drop the mapping that was inserted. - if (mat.getKind() == UnresolvedMaterialization::Target) { - for (Value input : op->getOperands()) - mapping.erase(input); - } - detachNestedAndErase(op); - } - - // Pop all of the newly created operations. - while (createdOps.size() != state.numCreatedOps) { - detachNestedAndErase(createdOps.back()); - createdOps.pop_back(); - } - // Pop all of the recorded ignored operations that are no longer valid. while (ignoredOps.size() != state.numIgnoredOperations) ignoredOps.pop_back(); - // Reset operations with changed results. - while (!operationsWithChangedResults.empty() && - operationsWithChangedResults.back() >= state.numReplacements) - operationsWithChangedResults.pop_back(); + while (eraseRewriter.erased.size() != state.numErased) + eraseRewriter.erased.pop_back(); } void ConversionPatternRewriterImpl::undoRewrites(unsigned numRewritesToKeep) { @@ -1413,8 +1219,7 @@ LogicalResult ConversionPatternRewriterImpl::remapValues( if (currentTypeConverter && desiredType && newOperandType != desiredType) { Location operandLoc = inputLoc ? *inputLoc : operand.getLoc(); Value castValue = buildUnresolvedTargetMaterialization( - operandLoc, newOperand, desiredType, currentTypeConverter, - unresolvedMaterializations); + operandLoc, newOperand, desiredType, currentTypeConverter); mapping.map(mapping.lookupOrDefault(newOperand), castValue); newOperand = castValue; } @@ -1425,7 +1230,8 @@ LogicalResult ConversionPatternRewriterImpl::remapValues( bool ConversionPatternRewriterImpl::isOpIgnored(Operation *op) const { // Check to see if this operation was replaced or its parent ignored. - return replacements.count(op) || ignoredOps.count(op->getParentOp()); + return ignoredOps.count(op->getParentOp()) || + hasRewrite<ReplaceOperationRewrite>(rewrites, op); } void ConversionPatternRewriterImpl::markNestedOpsIgnored(Operation *op) { @@ -1447,18 +1253,18 @@ void ConversionPatternRewriterImpl::markNestedOpsIgnored(Operation *op) { FailureOr<Block *> ConversionPatternRewriterImpl::convertBlockSignature( Block *block, const TypeConverter *converter, TypeConverter::SignatureConversion *conversion) { - FailureOr<Block *> result = - conversion ? argConverter.applySignatureConversion( - block, converter, *conversion, mapping, argReplacements) - : argConverter.convertSignature(block, converter, mapping, - argReplacements); - if (failed(result)) + if (conversion) + return applySignatureConversion(block, converter, *conversion); + + // If a converter wasn't provided, and the block wasn't already converted, + // there is nothing we can do. + if (!converter) return failure(); - if (Block *newBlock = *result) { - if (newBlock != block) - appendRewrite<BlockTypeConversionRewrite>(newBlock); - } - return result; + + // Try to convert the signature for the block with the provided converter. + if (auto conversion = converter->convertBlockSignature(block)) + return applySignatureConversion(block, converter, *conversion); + return failure(); } Block *ConversionPatternRewriterImpl::applySignatureConversion( @@ -1512,6 +1318,145 @@ LogicalResult ConversionPatternRewriterImpl::convertNonEntryRegionTypes( return success(); } +Block *ConversionPatternRewriterImpl::applySignatureConversion( + Block *block, const TypeConverter *converter, + TypeConverter::SignatureConversion &signatureConversion) { + MLIRContext *ctx = eraseRewriter.getContext(); + + // If no arguments are being changed or added, there is nothing to do. + unsigned origArgCount = block->getNumArguments(); + auto convertedTypes = signatureConversion.getConvertedTypes(); + if (llvm::equal(block->getArgumentTypes(), convertedTypes)) + return block; + + // Split the block at the beginning to get a new block to use for the updated + // signature. + Block *newBlock = block->splitBlock(block->begin()); + block->replaceAllUsesWith(newBlock); + // Unlink the block, but do not erase it yet, so that the change can be rolled + // back. + block->getParent()->getBlocks().remove(block); + + // Map all new arguments to the location of the argument they originate from. + SmallVector<Location> newLocs(convertedTypes.size(), + Builder(ctx).getUnknownLoc()); + for (unsigned i = 0; i < origArgCount; ++i) { + auto inputMap = signatureConversion.getInputMapping(i); + if (!inputMap || inputMap->replacementValue) + continue; + Location origLoc = block->getArgument(i).getLoc(); + for (unsigned j = 0; j < inputMap->size; ++j) + newLocs[inputMap->inputNo + j] = origLoc; + } + + SmallVector<Value, 4> newArgRange( + newBlock->addArguments(convertedTypes, newLocs)); + ArrayRef<Value> newArgs(newArgRange); + + // Remap each of the original arguments as determined by the signature + // conversion. + SmallVector<std::optional<ConvertedArgInfo>, 1> argInfo; + argInfo.resize(origArgCount); + + for (unsigned i = 0; i != origArgCount; ++i) { + auto inputMap = signatureConversion.getInputMapping(i); + if (!inputMap) + continue; + BlockArgument origArg = block->getArgument(i); + + // If inputMap->replacementValue is not nullptr, then the argument is + // dropped and a replacement value is provided to be the remappedValue. + if (inputMap->replacementValue) { + assert(inputMap->size == 0 && + "invalid to provide a replacement value when the argument isn't " + "dropped"); + mapping.map(origArg, inputMap->replacementValue); + appendRewrite<ReplaceBlockArgRewrite>(block, origArg); + continue; + } + + // Otherwise, this is a 1->1+ mapping. + auto replArgs = newArgs.slice(inputMap->inputNo, inputMap->size); + Value newArg; + + // If this is a 1->1 mapping and the types of new and replacement arguments + // match (i.e. it's an identity map), then the argument is mapped to its + // original type. + // FIXME: We simply pass through the replacement argument if there wasn't a + // converter, which isn't great as it allows implicit type conversions to + // appear. We should properly restructure this code to handle cases where a + // converter isn't provided and also to properly handle the case where an + // argument materialization is actually a temporary source materialization + // (e.g. in the case of 1->N). + if (replArgs.size() == 1 && + (!converter || replArgs[0].getType() == origArg.getType())) { + newArg = replArgs.front(); + } else { + Type origOutputType = origArg.getType(); + + // Legalize the argument output type. + Type outputType = origOutputType; + if (Type legalOutputType = converter->convertType(outputType)) + outputType = legalOutputType; + + newArg = buildUnresolvedArgumentMaterialization( + newBlock, origArg.getLoc(), replArgs, origOutputType, outputType, + converter); + } + + mapping.map(origArg, newArg); + appendRewrite<ReplaceBlockArgRewrite>(block, origArg); + argInfo[i] = ConvertedArgInfo(inputMap->inputNo, inputMap->size, newArg); + } + + appendRewrite<BlockTypeConversionRewrite>(newBlock, block, argInfo, + converter); + return newBlock; +} + +//===----------------------------------------------------------------------===// +// Materializations +//===----------------------------------------------------------------------===// + +/// Build an unresolved materialization operation given an output type and set +/// of input operands. +Value ConversionPatternRewriterImpl::buildUnresolvedMaterialization( + MaterializationKind kind, Block *insertBlock, Block::iterator insertPt, + Location loc, ValueRange inputs, Type outputType, Type origOutputType, + const TypeConverter *converter) { + // Avoid materializing an unnecessary cast. + if (inputs.size() == 1 && inputs.front().getType() == outputType) + return inputs.front(); + + // Create an unresolved materialization. We use a new OpBuilder to avoid + // tracking the materialization like we do for other operations. + OpBuilder builder(insertBlock, insertPt); + auto convertOp = + builder.create<UnrealizedConversionCastOp>(loc, outputType, inputs); + appendRewrite<UnresolvedMaterializationRewrite>(convertOp, converter, kind, + origOutputType); + return convertOp.getResult(0); +} +Value ConversionPatternRewriterImpl::buildUnresolvedArgumentMaterialization( + Block *block, Location loc, ValueRange inputs, Type origOutputType, + Type outputType, const TypeConverter *converter) { + return buildUnresolvedMaterialization(MaterializationKind::Argument, block, + block->begin(), loc, inputs, outputType, + origOutputType, converter); +} +Value ConversionPatternRewriterImpl::buildUnresolvedTargetMaterialization( + Location loc, Value input, Type outputType, + const TypeConverter *converter) { + Block *insertBlock = input.getParentBlock(); + Block::iterator insertPt = insertBlock->begin(); + if (OpResult inputRes = dyn_cast<OpResult>(input)) + insertPt = ++inputRes.getOwner()->getIterator(); + + return buildUnresolvedMaterialization(MaterializationKind::Target, + insertBlock, insertPt, loc, input, + outputType, outputType, converter); +} + //===----------------------------------------------------------------------===// // Rewriter Notification Hooks @@ -1523,7 +1468,7 @@ void ConversionPatternRewriterImpl::notifyOperationInserted( }); if (!previous.isSet()) { // This is a newly created op. - createdOps.push_back(op); + appendRewrite<CreateOperationRewrite>(op); return; } Operation *prevOp = previous.getPoint() == previous.getBlock()->end() @@ -1535,7 +1480,12 @@ void ConversionPatternRewriterImpl::notifyOperationInserted( void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op, ValueRange newValues) { assert(newValues.size() == op->getNumResults()); - assert(!replacements.count(op) && "operation was already replaced"); +#ifndef NDEBUG + for (auto &rewrite : rewrites) + if (auto *opReplacement = dyn_cast<ReplaceOperationRewrite>(rewrite.get())) + assert(opReplacement->getOperation() != op && + "operation was already replaced"); +#endif // NDEBUG // Track if any of the results changed, e.g. erased and replaced with null. bool resultChanged = false; @@ -1550,11 +1500,9 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op, mapping.map(result, newValue); resultChanged |= (newValue.getType() != result.getType()); } - if (resultChanged) - operationsWithChangedResults.push_back(replacements.size()); - // Record the requested operation replacement. - replacements.insert(std::make_pair(op, OpReplacement(currentTypeConverter))); + appendRewrite<ReplaceOperationRewrite>(op, currentTypeConverter, + resultChanged); // Mark this operation as recursively ignored so that we don't need to // convert any nested operations. @@ -1594,8 +1542,8 @@ void ConversionPatternRewriterImpl::notifyMatchFailure( Diagnostic diag(loc, DiagnosticSeverity::Remark); reasonCallback(diag); logger.startLine() << "** Failure : " << diag.str() << "\n"; - if (notifyCallback) - notifyCallback(diag); + if (config.notifyCallback) + config.notifyCallback(diag); }); } @@ -1603,9 +1551,10 @@ void ConversionPatternRewriterImpl::notifyMatchFailure( // ConversionPatternRewriter //===----------------------------------------------------------------------===// -ConversionPatternRewriter::ConversionPatternRewriter(MLIRContext *ctx) +ConversionPatternRewriter::ConversionPatternRewriter( + MLIRContext *ctx, const ConversionConfig &config) : PatternRewriter(ctx), - impl(new detail::ConversionPatternRewriterImpl(*this)) { + impl(new detail::ConversionPatternRewriterImpl(ctx, config)) { setListener(impl.get()); } @@ -1649,8 +1598,6 @@ void ConversionPatternRewriter::eraseOp(Operation *op) { } void ConversionPatternRewriter::eraseBlock(Block *block) { - impl->notifyBlockIsBeingErased(block); - // Mark all ops for erasure. for (Operation &op : *block) eraseOp(&op); @@ -1659,6 +1606,7 @@ void ConversionPatternRewriter::eraseBlock(Block *block) { // object and will be actually destroyed when rewrites are applied. This // allows us to keep the operations in the block live and undo the removal by // re-inserting the block. + impl->notifyBlockIsBeingErased(block); block->getParent()->getBlocks().remove(block); } @@ -1688,7 +1636,7 @@ void ConversionPatternRewriter::replaceUsesOfBlockArgument(BlockArgument from, << "'(in region of '" << parentOp->getName() << "'(" << from.getOwner()->getParentOp() << ")\n"; }); - impl->argReplacements.push_back(from); + impl->appendRewrite<ReplaceBlockArgRewrite>(from.getOwner(), from); impl->mapping.map(impl->mapping.lookupOrDefault(from), to); } @@ -2022,13 +1970,16 @@ OperationLegalizer::legalizeWithFold(Operation *op, rewriter.replaceOp(op, replacementValues); // Recursively legalize any new constant operations. - for (unsigned i = curState.numCreatedOps, e = rewriterImpl.createdOps.size(); + for (unsigned i = curState.numRewrites, e = rewriterImpl.rewrites.size(); i != e; ++i) { - Operation *cstOp = rewriterImpl.createdOps[i]; - if (failed(legalize(cstOp, rewriter))) { + auto *createOp = + dyn_cast<CreateOperationRewrite>(rewriterImpl.rewrites[i].get()); + if (!createOp) + continue; + if (failed(legalize(createOp->getOperation(), rewriter))) { LLVM_DEBUG(logFailure(rewriterImpl.logger, "failed to legalize generated constant '{0}'", - cstOp->getName())); + createOp->getOperation()->getName())); rewriterImpl.resetState(curState); return failure(); } @@ -2054,12 +2005,12 @@ OperationLegalizer::legalizeWithPattern(Operation *op, assert(rewriterImpl.pendingRootUpdates.empty() && "dangling root updates"); LLVM_DEBUG({ logFailure(rewriterImpl.logger, "pattern failed to match"); - if (rewriterImpl.notifyCallback) { + if (rewriterImpl.config.notifyCallback) { Diagnostic diag(op->getLoc(), DiagnosticSeverity::Remark); diag << "Failed to apply pattern \"" << pattern.getDebugName() << "\" on op:\n" << *op; - rewriterImpl.notifyCallback(diag); + rewriterImpl.config.notifyCallback(diag); } }); rewriterImpl.resetState(curState); @@ -2112,16 +2063,13 @@ OperationLegalizer::legalizePatternResult(Operation *op, const Pattern &pattern, #ifndef NDEBUG assert(impl.pendingRootUpdates.empty() && "dangling root updates"); - // Check that the root was either replaced or updated in place. + auto newRewrites = llvm::drop_begin(impl.rewrites, curState.numRewrites); auto replacedRoot = [&] { - return llvm::any_of( - llvm::drop_begin(impl.replacements, curState.numReplacements), - [op](auto &it) { return it.first == op; }); + return hasRewrite<ReplaceOperationRewrite>(newRewrites, op); }; auto updatedRootInPlace = [&] { - return hasRewrite<ModifyOperationRewrite>( - llvm::drop_begin(impl.rewrites, curState.numRewrites), op); + return hasRewrite<ModifyOperationRewrite>(newRewrites, op); }; assert((replacedRoot() || updatedRootInPlace()) && "expected pattern to replace the root operation"); @@ -2154,7 +2102,8 @@ LogicalResult OperationLegalizer::legalizePatternBlockRewrites( if (!rewrite) continue; Block *block = rewrite->getBlock(); - if (isa<BlockTypeConversionRewrite, EraseBlockRewrite>(rewrite)) + if (isa<BlockTypeConversionRewrite, EraseBlockRewrite, + ReplaceBlockArgRewrite>(rewrite)) continue; // Only check blocks outside of the current operation. Operation *parentOp = block->getParentOp(); @@ -2177,9 +2126,14 @@ LogicalResult OperationLegalizer::legalizePatternBlockRewrites( // blocks in regions created by this pattern will already be legalized later // on. If we haven't built the set yet, build it now. if (operationsToIgnore.empty()) { - auto createdOps = ArrayRef<Operation *>(impl.createdOps) - .drop_front(state.numCreatedOps); - operationsToIgnore.insert(createdOps.begin(), createdOps.end()); + for (unsigned i = state.numRewrites, e = impl.rewrites.size(); i != e; + ++i) { + auto *createOp = + dyn_cast<CreateOperationRewrite>(impl.rewrites[i].get()); + if (!createOp) + continue; + operationsToIgnore.insert(createOp->getOperation()); + } } // If this operation should be considered for re-legalization, try it. @@ -2197,8 +2151,11 @@ LogicalResult OperationLegalizer::legalizePatternBlockRewrites( LogicalResult OperationLegalizer::legalizePatternCreatedOperations( ConversionPatternRewriter &rewriter, ConversionPatternRewriterImpl &impl, RewriterState &state, RewriterState &newState) { - for (int i = state.numCreatedOps, e = newState.numCreatedOps; i != e; ++i) { - Operation *op = impl.createdOps[i]; + for (int i = state.numRewrites, e = newState.numRewrites; i != e; ++i) { + auto *createOp = dyn_cast<CreateOperationRewrite>(impl.rewrites[i].get()); + if (!createOp) + continue; + Operation *op = createOp->getOperation(); if (failed(legalize(op, rewriter))) { LLVM_DEBUG(logFailure(impl.logger, "failed to legalize generated operation '{0}'({1})", @@ -2432,21 +2389,21 @@ enum OpConversionMode { /// applied to the operations on success. Analysis, }; +} // namespace +namespace mlir { // This class converts operations to a given conversion target via a set of // rewrite patterns. The conversion behaves differently depending on the // conversion mode. struct OperationConverter { explicit OperationConverter(const ConversionTarget &target, const FrozenRewritePatternSet &patterns, - OpConversionMode mode, - DenseSet<Operation *> *trackedOps = nullptr) - : opLegalizer(target, patterns), mode(mode), trackedOps(trackedOps) {} + const ConversionConfig &config, + OpConversionMode mode) + : opLegalizer(target, patterns), config(config), mode(mode) {} /// Converts the given operations to the conversion target. - LogicalResult - convertOperations(ArrayRef<Operation *> ops, - function_ref<void(Diagnostic &)> notifyCallback = nullptr); + LogicalResult convertOperations(ArrayRef<Operation *> ops); private: /// Converts an operation with the given rewriter. @@ -2483,16 +2440,13 @@ private: /// The legalizer to use when converting operations. OperationLegalizer opLegalizer; + /// Dialect conversion configuration. + ConversionConfig config; + /// The conversion mode to use when legalizing operations. OpConversionMode mode; - - /// A set of pre-existing operations. When mode == OpConversionMode::Analysis, - /// this is populated with ops found to be legalizable to the target. - /// When mode == OpConversionMode::Partial, this is populated with ops found - /// *not* to be legalizable to the target. - DenseSet<Operation *> *trackedOps; }; -} // namespace +} // namespace mlir LogicalResult OperationConverter::convert(ConversionPatternRewriter &rewriter, Operation *op) { @@ -2504,28 +2458,27 @@ LogicalResult OperationConverter::convert(ConversionPatternRewriter &rewriter, return op->emitError() << "failed to legalize operation '" << op->getName() << "'"; // Partial conversions allow conversions to fail iff the operation was not - // explicitly marked as illegal. If the user provided a nonlegalizableOps - // set, non-legalizable ops are included. + // explicitly marked as illegal. If the user provided a `unlegalizedOps` + // set, non-legalizable ops are added to that set. if (mode == OpConversionMode::Partial) { if (opLegalizer.isIllegal(op)) return op->emitError() << "failed to legalize operation '" << op->getName() << "' that was explicitly marked illegal"; - if (trackedOps) - trackedOps->insert(op); + if (config.unlegalizedOps) + config.unlegalizedOps->insert(op); } } else if (mode == OpConversionMode::Analysis) { // Analysis conversions don't fail if any operations fail to legalize, // they are only interested in the operations that were successfully // legalized. - trackedOps->insert(op); + if (config.legalizableOps) + config.legalizableOps->insert(op); } return success(); } -LogicalResult OperationConverter::convertOperations( - ArrayRef<Operation *> ops, - function_ref<void(Diagnostic &)> notifyCallback) { +LogicalResult OperationConverter::convertOperations(ArrayRef<Operation *> ops) { if (ops.empty()) return success(); const ConversionTarget &target = opLegalizer.getTarget(); @@ -2546,33 +2499,25 @@ LogicalResult OperationConverter::convertOperations( } // Convert each operation and discard rewrites on failure. - ConversionPatternRewriter rewriter(ops.front()->getContext()); + ConversionPatternRewriter rewriter(ops.front()->getContext(), config); ConversionPatternRewriterImpl &rewriterImpl = rewriter.getImpl(); - rewriterImpl.notifyCallback = notifyCallback; for (auto *op : toConvert) if (failed(convert(rewriter, op))) - return rewriterImpl.discardRewrites(), failure(); + return rewriterImpl.undoRewrites(), failure(); // Now that all of the operations have been converted, finalize the conversion // process to ensure any lingering conversion artifacts are cleaned up and // legalized. if (failed(finalize(rewriter))) - return rewriterImpl.discardRewrites(), failure(); + return rewriterImpl.undoRewrites(), failure(); // After a successful conversion, apply rewrites if this is not an analysis // conversion. if (mode == OpConversionMode::Analysis) { - rewriterImpl.discardRewrites(); + rewriterImpl.undoRewrites(); } else { rewriterImpl.applyRewrites(); - - // It is possible for a later pattern to erase an op that was originally - // identified as illegal and added to the trackedOps, remove it now after - // replacements have been computed. - if (trackedOps) - for (auto &repl : rewriterImpl.replacements) - trackedOps->erase(repl.first); } return success(); } @@ -2586,21 +2531,20 @@ OperationConverter::finalize(ConversionPatternRewriter &rewriter) { failed(legalizeConvertedArgumentTypes(rewriter, rewriterImpl))) return failure(); - if (rewriterImpl.operationsWithChangedResults.empty()) - return success(); - // Process requested operation replacements. - for (unsigned i = 0, e = rewriterImpl.operationsWithChangedResults.size(); - i != e; ++i) { - unsigned replIdx = rewriterImpl.operationsWithChangedResults[i]; - auto &repl = *(rewriterImpl.replacements.begin() + replIdx); - for (OpResult result : repl.first->getResults()) { + for (unsigned i = 0; i < rewriterImpl.rewrites.size(); ++i) { + auto *opReplacement = + dyn_cast<ReplaceOperationRewrite>(rewriterImpl.rewrites[i].get()); + if (!opReplacement || !opReplacement->hasChangedResults()) + continue; + Operation *op = opReplacement->getOperation(); + for (OpResult result : op->getResults()) { Value newValue = rewriterImpl.mapping.lookupOrNull(result); // If the operation result was replaced with null, all of the uses of this // value should be replaced. if (!newValue) { - if (failed(legalizeErasedResult(repl.first, result, rewriterImpl))) + if (failed(legalizeErasedResult(op, result, rewriterImpl))) return failure(); continue; } @@ -2614,15 +2558,11 @@ OperationConverter::finalize(ConversionPatternRewriter &rewriter) { inverseMapping = rewriterImpl.mapping.getInverse(); // Legalize this result. - rewriter.setInsertionPoint(repl.first); - if (failed(legalizeChangedResultType(repl.first, result, newValue, - repl.second.converter, rewriter, - rewriterImpl, *inverseMapping))) + rewriter.setInsertionPoint(op); + if (failed(legalizeChangedResultType( + op, result, newValue, opReplacement->getConverter(), rewriter, + rewriterImpl, *inverseMapping))) return failure(); - - // Update the end iterator for this loop in the case it was updated - // when legalizing generated conversion operations. - e = rewriterImpl.operationsWithChangedResults.size(); } } return success(); @@ -2639,8 +2579,17 @@ LogicalResult OperationConverter::legalizeConvertedArgumentTypes( }); return liveUserIt == val.user_end() ? nullptr : *liveUserIt; }; - return rewriterImpl.argConverter.materializeLiveConversions( - rewriterImpl.mapping, rewriter, findLiveUser); + // Note: `rewrites` may be reallocated as the loop is running. + for (int64_t i = 0; i < static_cast<int64_t>(rewriterImpl.rewrites.size()); + ++i) { + auto &rewrite = rewriterImpl.rewrites[i]; + if (auto *blockTypeConversionRewrite = + dyn_cast<BlockTypeConversionRewrite>(rewrite.get())) + if (failed(blockTypeConversionRewrite->materializeLiveConversions( + findLiveUser))) + return failure(); + } + return success(); } /// Replace the results of a materialization operation with the given values. @@ -2672,11 +2621,12 @@ replaceMaterialization(ConversionPatternRewriterImpl &rewriterImpl, /// Compute all of the unresolved materializations that will persist beyond the /// conversion process, and require inserting a proper user materialization for. static void computeNecessaryMaterializations( - DenseMap<Operation *, UnresolvedMaterialization *> &materializationOps, + DenseMap<Operation *, UnresolvedMaterializationRewrite *> + &materializationOps, ConversionPatternRewriter &rewriter, ConversionPatternRewriterImpl &rewriterImpl, DenseMap<Value, SmallVector<Value>> &inverseMapping, - SetVector<UnresolvedMaterialization *> &necessaryMaterializations) { + SetVector<UnresolvedMaterializationRewrite *> &necessaryMaterializations) { auto isLive = [&](Value value) { auto findFn = [&](Operation *user) { auto matIt = materializationOps.find(user); @@ -2711,14 +2661,17 @@ static void computeNecessaryMaterializations( return Value(); }; - SetVector<UnresolvedMaterialization *> worklist; - for (auto &mat : rewriterImpl.unresolvedMaterializations) { - materializationOps.try_emplace(mat.getOp(), &mat); - worklist.insert(&mat); + SetVector<UnresolvedMaterializationRewrite *> worklist; + for (auto &rewrite : rewriterImpl.rewrites) { + auto *mat = dyn_cast<UnresolvedMaterializationRewrite>(rewrite.get()); + if (!mat) + continue; + materializationOps.try_emplace(mat->getOperation(), mat); + worklist.insert(mat); } while (!worklist.empty()) { - UnresolvedMaterialization *mat = worklist.pop_back_val(); - UnrealizedConversionCastOp op = mat->getOp(); + UnresolvedMaterializationRewrite *mat = worklist.pop_back_val(); + UnrealizedConversionCastOp op = mat->getOperation(); // We currently only handle target materializations here. assert(op->getNumResults() == 1 && "unexpected materialization type"); @@ -2760,7 +2713,7 @@ static void computeNecessaryMaterializations( auto isBlockArg = [](Value v) { return isa<BlockArgument>(v); }; if (llvm::any_of(op->getOperands(), isBlockArg) || llvm::any_of(inverseMapping[op->getResult(0)], isBlockArg)) { - mat->setKind(UnresolvedMaterialization::Argument); + mat->setMaterializationKind(MaterializationKind::Argument); } // If the materialization does not have any live users, we don't need to @@ -2770,7 +2723,7 @@ static void computeNecessaryMaterializations( // value replacement even if the types differ in some cases. When those // patterns are fixed, we can drop the argument special case here. bool isMaterializationLive = isLive(opResult); - if (mat->getKind() == UnresolvedMaterialization::Argument) + if (mat->getMaterializationKind() == MaterializationKind::Argument) isMaterializationLive |= llvm::any_of(inverseMapping[opResult], isLive); if (!isMaterializationLive) continue; @@ -2790,8 +2743,9 @@ static void computeNecessaryMaterializations( /// Legalize the given unresolved materialization. Returns success if the /// materialization was legalized, failure otherise. static LogicalResult legalizeUnresolvedMaterialization( - UnresolvedMaterialization &mat, - DenseMap<Operation *, UnresolvedMaterialization *> &materializationOps, + UnresolvedMaterializationRewrite &mat, + DenseMap<Operation *, UnresolvedMaterializationRewrite *> + &materializationOps, ConversionPatternRewriter &rewriter, ConversionPatternRewriterImpl &rewriterImpl, DenseMap<Value, SmallVector<Value>> &inverseMapping) { @@ -2811,7 +2765,7 @@ static LogicalResult legalizeUnresolvedMaterialization( return Value(); }; - UnrealizedConversionCastOp op = mat.getOp(); + UnrealizedConversionCastOp op = mat.getOperation(); if (!rewriterImpl.ignoredOps.insert(op)) return success(); @@ -2861,8 +2815,8 @@ static LogicalResult legalizeUnresolvedMaterialization( rewriter.setInsertionPoint(op); Value newMaterialization; - switch (mat.getKind()) { - case UnresolvedMaterialization::Argument: + switch (mat.getMaterializationKind()) { + case MaterializationKind::Argument: // Try to materialize an argument conversion. // FIXME: The current argument materialization hook expects the original // output type, even though it doesn't use that as the actual output type @@ -2879,7 +2833,7 @@ static LogicalResult legalizeUnresolvedMaterialization( // If an argument materialization failed, fallback to trying a target // materialization. [[fallthrough]]; - case UnresolvedMaterialization::Target: + case MaterializationKind::Target: newMaterialization = converter->materializeTargetConversion( rewriter, op->getLoc(), outputType, inputOperands); break; @@ -2907,14 +2861,12 @@ LogicalResult OperationConverter::legalizeUnresolvedMaterializations( ConversionPatternRewriter &rewriter, ConversionPatternRewriterImpl &rewriterImpl, std::optional<DenseMap<Value, SmallVector<Value>>> &inverseMapping) { - if (rewriterImpl.unresolvedMaterializations.empty()) - return success(); inverseMapping = rewriterImpl.mapping.getInverse(); // As an initial step, compute all of the inserted materializations that we // expect to persist beyond the conversion process. - DenseMap<Operation *, UnresolvedMaterialization *> materializationOps; - SetVector<UnresolvedMaterialization *> necessaryMaterializations; + DenseMap<Operation *, UnresolvedMaterializationRewrite *> materializationOps; + SetVector<UnresolvedMaterializationRewrite *> necessaryMaterializations; computeNecessaryMaterializations(materializationOps, rewriter, rewriterImpl, *inverseMapping, necessaryMaterializations); @@ -3535,57 +3487,51 @@ void mlir::registerConversionPDLFunctions(RewritePatternSet &patterns) { //===----------------------------------------------------------------------===// // Partial Conversion -LogicalResult -mlir::applyPartialConversion(ArrayRef<Operation *> ops, - const ConversionTarget &target, - const FrozenRewritePatternSet &patterns, - DenseSet<Operation *> *unconvertedOps) { - OperationConverter opConverter(target, patterns, OpConversionMode::Partial, - unconvertedOps); +LogicalResult mlir::applyPartialConversion( + ArrayRef<Operation *> ops, const ConversionTarget &target, + const FrozenRewritePatternSet &patterns, ConversionConfig config) { + OperationConverter opConverter(target, patterns, config, + OpConversionMode::Partial); return opConverter.convertOperations(ops); } LogicalResult mlir::applyPartialConversion(Operation *op, const ConversionTarget &target, const FrozenRewritePatternSet &patterns, - DenseSet<Operation *> *unconvertedOps) { - return applyPartialConversion(llvm::ArrayRef(op), target, patterns, - unconvertedOps); + ConversionConfig config) { + return applyPartialConversion(llvm::ArrayRef(op), target, patterns, config); } //===----------------------------------------------------------------------===// // Full Conversion -LogicalResult -mlir::applyFullConversion(ArrayRef<Operation *> ops, - const ConversionTarget &target, - const FrozenRewritePatternSet &patterns) { - OperationConverter opConverter(target, patterns, OpConversionMode::Full); +LogicalResult mlir::applyFullConversion(ArrayRef<Operation *> ops, + const ConversionTarget &target, + const FrozenRewritePatternSet &patterns, + ConversionConfig config) { + OperationConverter opConverter(target, patterns, config, + OpConversionMode::Full); return opConverter.convertOperations(ops); } -LogicalResult -mlir::applyFullConversion(Operation *op, const ConversionTarget &target, - const FrozenRewritePatternSet &patterns) { - return applyFullConversion(llvm::ArrayRef(op), target, patterns); +LogicalResult mlir::applyFullConversion(Operation *op, + const ConversionTarget &target, + const FrozenRewritePatternSet &patterns, + ConversionConfig config) { + return applyFullConversion(llvm::ArrayRef(op), target, patterns, config); } //===----------------------------------------------------------------------===// // Analysis Conversion -LogicalResult -mlir::applyAnalysisConversion(ArrayRef<Operation *> ops, - ConversionTarget &target, - const FrozenRewritePatternSet &patterns, - DenseSet<Operation *> &convertedOps, - function_ref<void(Diagnostic &)> notifyCallback) { - OperationConverter opConverter(target, patterns, OpConversionMode::Analysis, - &convertedOps); - return opConverter.convertOperations(ops, notifyCallback); +LogicalResult mlir::applyAnalysisConversion( + ArrayRef<Operation *> ops, ConversionTarget &target, + const FrozenRewritePatternSet &patterns, ConversionConfig config) { + OperationConverter opConverter(target, patterns, config, + OpConversionMode::Analysis); + return opConverter.convertOperations(ops); } LogicalResult mlir::applyAnalysisConversion(Operation *op, ConversionTarget &target, const FrozenRewritePatternSet &patterns, - DenseSet<Operation *> &convertedOps, - function_ref<void(Diagnostic &)> notifyCallback) { - return applyAnalysisConversion(llvm::ArrayRef(op), target, patterns, - convertedOps, notifyCallback); + ConversionConfig config) { + return applyAnalysisConversion(llvm::ArrayRef(op), target, patterns, config); } |