267 lines · cpp
1//===- SparsificationAndBufferizationPass.cpp - Tensor to Memref Lowering -===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"10 11#include "mlir/Dialect/Affine/IR/AffineOps.h"12#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"13#include "mlir/Dialect/Bufferization/IR/Bufferization.h"14#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"15#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"16#include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"17#include "mlir/Dialect/Bufferization/Transforms/Passes.h"18#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"19#include "mlir/Dialect/Func/IR/FuncOps.h"20#include "mlir/Dialect/GPU/IR/GPUDialect.h"21#include "mlir/Dialect/LLVMIR/LLVMDialect.h"22#include "mlir/Dialect/Linalg/IR/Linalg.h"23#include "mlir/Dialect/SCF/IR/SCF.h"24#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"25#include "mlir/Dialect/Vector/IR/VectorOps.h"26#include "mlir/Pass/PassManager.h"27#include "mlir/Transforms/Passes.h"28 29using namespace mlir;30 31namespace mlir {32 33#define GEN_PASS_DEF_SPARSIFICATIONANDBUFFERIZATION34#include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"35 36namespace sparse_tensor {37 38/// Return `true` if one of the given types is a sparse tensor type.39static bool containsSparseTensor(TypeRange types) {40 for (Type t : types)41 if (isa<TensorType>(t) && getSparseTensorEncoding(t))42 return true;43 return false;44}45 46/// A pass that lowers tensor ops to memref ops, regardless of whether they are47/// dense or sparse.48///49/// One-Shot Analysis is used to detect RaW conflicts and to insert buffer50/// copies of the tensor level (`insertTensorCopies`). Afterwards, the lowering51/// of tensor ops to memref ops follows a different code path depending on52/// whether the op is sparse or dense:53///54/// * Sparse tensor ops are lowered through Sparsification and follow-up pass55/// that lowers sparse_tensor dialect ops.56/// * Dense tensor ops are lowered through BufferizableOpInterface57/// implementations.58class SparsificationAndBufferizationPass59 : public impl::SparsificationAndBufferizationBase<60 SparsificationAndBufferizationPass> {61public:62 // Private pass options only.63 SparsificationAndBufferizationPass(64 const bufferization::OneShotBufferizationOptions &bufferizationOptions,65 const SparsificationOptions &sparsificationOptions,66 bool createSparseDeallocs, bool enableRuntimeLibrary,67 bool enableBufferInitialization)68 : bufferizationOptions(bufferizationOptions),69 sparsificationOptions(sparsificationOptions),70 createSparseDeallocs(createSparseDeallocs),71 enableRuntimeLibrary(enableRuntimeLibrary),72 enableBufferInitialization(enableBufferInitialization) {}73 // Private pass options and visible pass options.74 SparsificationAndBufferizationPass(75 const bufferization::OneShotBufferizationOptions &bufferizationOptions,76 const SparsificationOptions &sparsificationOptions,77 bool createSparseDeallocs, bool enableRuntimeLibrary,78 bool enableBufferInitialization, unsigned vl, bool vla, bool index32,79 bool gpu, SparseEmitStrategy emitStrategy,80 SparseParallelizationStrategy parallelizationStrategy)81 : bufferizationOptions(bufferizationOptions),82 sparsificationOptions(sparsificationOptions),83 createSparseDeallocs(createSparseDeallocs),84 enableRuntimeLibrary(enableRuntimeLibrary),85 enableBufferInitialization(enableBufferInitialization) {86 // Set the visible pass options explicitly.87 vectorLength = vl;88 enableVLAVectorization = vla;89 enableSIMDIndex32 = index32;90 enableGPULibgen = gpu;91 sparseEmitStrategy = emitStrategy;92 parallelization = parallelizationStrategy;93 }94 95 /// Bufferize all dense ops. This assumes that no further analysis is needed96 /// and that all required buffer copies were already inserted by97 /// `insertTensorCopies` in the form of `bufferization.alloc_tensor` ops.98 LogicalResult runDenseBufferization() {99 bufferization::OneShotBufferizationOptions updatedOptions =100 bufferizationOptions;101 // Skip all sparse ops.102 updatedOptions.opFilter.denyOperation([&](Operation *op) {103 if (containsSparseTensor(TypeRange(op->getResults())) ||104 containsSparseTensor(TypeRange(op->getOperands())))105 return true;106 if (auto funcOp = dyn_cast<func::FuncOp>(op)) {107 FunctionType funcType = funcOp.getFunctionType();108 if (containsSparseTensor(funcType.getInputs()) ||109 containsSparseTensor(funcType.getResults()))110 return true;111 }112 return false;113 });114 115 bufferization::BufferizationState bufferizationState;116 117 if (failed(bufferization::bufferizeModuleOp(getOperation(), updatedOptions,118 bufferizationState)))119 return failure();120 121 bufferization::removeBufferizationAttributesInModule(getOperation());122 return success();123 }124 125 void runOnOperation() override {126 // Overrides the default emit strategy using user-provided value.127 this->sparsificationOptions.sparseEmitStrategy = sparseEmitStrategy;128 129 // Overrides the default parallelization strategy using user-provided value.130 this->sparsificationOptions.parallelizationStrategy = parallelization;131 132 // Run enabling transformations.133 {134 OpPassManager pm("builtin.module");135 pm.addPass(createPreSparsificationRewritePass());136 pm.addNestedPass<func::FuncOp>(137 bufferization::createEmptyTensorToAllocTensorPass());138 if (failed(runPipeline(pm, getOperation())))139 return signalPassFailure();140 }141 142 // Insert tensor copies. This step runs One-Shot Analysis (which analyzes143 // SSA use-def chains of tensor IR) and decides where buffer copies are144 // needed and where buffers can be written to in-place. These decisions are145 // materialized in the IR in the form of `bufferization.alloc_tensor` ops.146 //147 // Note: All following steps in this pass must be careful not to modify the148 // structure of the IR (i.e., tensor use-def chains), as that could149 // invalidate the results of the analysis. From now on, only small and150 // localized rewrites are allowed, such as replacing a tensor op with its151 // memref equivalent.152 bufferization::BufferizationState bufferizationState;153 154 if (failed(bufferization::insertTensorCopies(155 getOperation(), bufferizationOptions, bufferizationState)))156 return signalPassFailure();157 158 // Option `testAnalysisOnly` is a debug/testing flag. If set, the results of159 // OneShotAnalysis are added to the IR via attributes. In that case, do not160 // continue with the remaining pipeline.161 if (bufferizationOptions.testAnalysisOnly)162 return;163 164 // Bufferize all sparse ops. No further analysis is needed. All required165 // buffer copies were already inserted by `insertTensorCopies` in the form166 // of `bufferization.alloc_tensor` ops.167 {168 OpPassManager pm("builtin.module");169 if (enableGPULibgen)170 pm.addPass(createSparseGPUCodegenPass(0, enableRuntimeLibrary));171 pm.addPass(createSparseReinterpretMapPass(ReinterpretMapScope::kAll));172 pm.addPass(createSparsificationPass(sparsificationOptions));173 if (sparsificationOptions.sparseEmitStrategy ==174 SparseEmitStrategy::kSparseIterator) {175 pm.addNestedPass<func::FuncOp>(createSparseSpaceCollapsePass());176 pm.addNestedPass<func::FuncOp>(createLowerSparseIterationToSCFPass());177 }178 179 pm.addNestedPass<func::FuncOp>(createStageSparseOperationsPass());180 pm.addPass(createLowerSparseOpsToForeachPass(enableRuntimeLibrary,181 /*enableConvert=*/true));182 pm.addPass(183 createSparseReinterpretMapPass(ReinterpretMapScope::kExceptGeneric));184 pm.addNestedPass<func::FuncOp>(createLowerForeachToSCFPass());185 pm.addPass(mlir::createLoopInvariantCodeMotionPass());186 if (vectorLength > 0) {187 pm.addPass(createSparseVectorizationPass(188 vectorLength, enableVLAVectorization, enableSIMDIndex32));189 }190 if (enableRuntimeLibrary) {191 pm.addPass(createSparseTensorConversionPass());192 } else {193 pm.addPass(createSparseTensorCodegenPass(createSparseDeallocs,194 enableBufferInitialization));195 pm.addPass(createSparseBufferRewritePass(enableBufferInitialization));196 }197 if (failed(runPipeline(pm, getOperation())))198 return signalPassFailure();199 }200 201 // Bufferize all dense ops.202 if (failed(runDenseBufferization()))203 signalPassFailure();204 }205 206private:207 bufferization::OneShotBufferizationOptions bufferizationOptions;208 SparsificationOptions sparsificationOptions;209 bool createSparseDeallocs;210 bool enableRuntimeLibrary;211 bool enableBufferInitialization;212};213 214} // namespace sparse_tensor215} // namespace mlir216 217mlir::bufferization::OneShotBufferizationOptions218mlir::getBufferizationOptionsForSparsification(bool analysisOnly) {219 using namespace mlir::bufferization;220 OneShotBufferizationOptions options;221 options.bufferizeFunctionBoundaries = true;222 options.setFunctionBoundaryTypeConversion(LayoutMapOption::IdentityLayoutMap);223 options.unknownTypeConverterFn = [](TensorType tensorType,224 Attribute memorySpace,225 const BufferizationOptions &options) {226 return getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace);227 };228 if (analysisOnly) {229 options.testAnalysisOnly = true;230 options.printConflicts = true;231 }232 // Since this mini-pipeline may be used in alternative pipelines (viz.233 // different from the default "sparsifier" pipeline) where unknown ops234 // are handled by alternative bufferization methods that are downstream235 // of this mini-pipeline, we allow unknown ops by default (failure to236 // bufferize is eventually apparent by failing to convert to LLVM IR).237 options.allowUnknownOps = true;238 return options;239}240 241std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass() {242 SparsificationOptions sparseOptions;243 return std::make_unique<244 mlir::sparse_tensor::SparsificationAndBufferizationPass>(245 getBufferizationOptionsForSparsification(/*analysisOnly=*/false),246 sparseOptions,247 /*createSparseDeallocs=*/false,248 /*enableRuntimeLibrary=*/false,249 /*enableBufferInitialization=*/false);250}251 252std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass(253 const bufferization::OneShotBufferizationOptions &bufferizationOptions,254 const SparsificationOptions &sparsificationOptions,255 bool createSparseDeallocs, bool enableRuntimeLibrary,256 bool enableBufferInitialization, unsigned vectorLength,257 bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen,258 SparseEmitStrategy emitStrategy,259 SparseParallelizationStrategy parallelizationStrategy) {260 return std::make_unique<261 mlir::sparse_tensor::SparsificationAndBufferizationPass>(262 bufferizationOptions, sparsificationOptions, createSparseDeallocs,263 enableRuntimeLibrary, enableBufferInitialization, vectorLength,264 enableVLAVectorization, enableSIMDIndex32, enableGPULibgen, emitStrategy,265 parallelizationStrategy);266}267