MLIR  22.0.0git
SparsificationAndBufferizationPass.cpp
Go to the documentation of this file.
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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
10 
27 #include "mlir/Pass/PassManager.h"
28 #include "mlir/Transforms/Passes.h"
29 
30 using namespace mlir;
31 
32 namespace mlir {
33 
34 #define GEN_PASS_DEF_SPARSIFICATIONANDBUFFERIZATION
35 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
36 
37 namespace sparse_tensor {
38 
39 /// Return `true` if one of the given types is a sparse tensor type.
40 static bool containsSparseTensor(TypeRange types) {
41  for (Type t : types)
42  if (isa<TensorType>(t) && getSparseTensorEncoding(t))
43  return true;
44  return false;
45 }
46 
47 /// A pass that lowers tensor ops to memref ops, regardless of whether they are
48 /// dense or sparse.
49 ///
50 /// One-Shot Analysis is used to detect RaW conflicts and to insert buffer
51 /// copies of the tensor level (`insertTensorCopies`). Afterwards, the lowering
52 /// of tensor ops to memref ops follows a different code path depending on
53 /// whether the op is sparse or dense:
54 ///
55 /// * Sparse tensor ops are lowered through Sparsification and follow-up pass
56 /// that lowers sparse_tensor dialect ops.
57 /// * Dense tensor ops are lowered through BufferizableOpInterface
58 /// implementations.
60  : public impl::SparsificationAndBufferizationBase<
61  SparsificationAndBufferizationPass> {
62 public:
63  // Private pass options only.
65  const bufferization::OneShotBufferizationOptions &bufferizationOptions,
66  const SparsificationOptions &sparsificationOptions,
67  bool createSparseDeallocs, bool enableRuntimeLibrary,
68  bool enableBufferInitialization)
69  : bufferizationOptions(bufferizationOptions),
70  sparsificationOptions(sparsificationOptions),
71  createSparseDeallocs(createSparseDeallocs),
72  enableRuntimeLibrary(enableRuntimeLibrary),
73  enableBufferInitialization(enableBufferInitialization) {}
74  // Private pass options and visible pass options.
76  const bufferization::OneShotBufferizationOptions &bufferizationOptions,
77  const SparsificationOptions &sparsificationOptions,
78  bool createSparseDeallocs, bool enableRuntimeLibrary,
79  bool enableBufferInitialization, unsigned vl, bool vla, bool index32,
80  bool gpu, SparseEmitStrategy emitStrategy,
81  SparseParallelizationStrategy parallelizationStrategy)
82  : bufferizationOptions(bufferizationOptions),
83  sparsificationOptions(sparsificationOptions),
84  createSparseDeallocs(createSparseDeallocs),
85  enableRuntimeLibrary(enableRuntimeLibrary),
86  enableBufferInitialization(enableBufferInitialization) {
87  // Set the visible pass options explicitly.
88  vectorLength = vl;
89  enableVLAVectorization = vla;
90  enableSIMDIndex32 = index32;
91  enableGPULibgen = gpu;
92  sparseEmitStrategy = emitStrategy;
93  parallelization = parallelizationStrategy;
94  }
95 
96  /// Bufferize all dense ops. This assumes that no further analysis is needed
97  /// and that all required buffer copies were already inserted by
98  /// `insertTensorCopies` in the form of `bufferization.alloc_tensor` ops.
99  LogicalResult runDenseBufferization() {
101  bufferizationOptions;
102  // Skip all sparse ops.
103  updatedOptions.opFilter.denyOperation([&](Operation *op) {
106  return true;
107  if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
108  FunctionType funcType = funcOp.getFunctionType();
109  if (containsSparseTensor(funcType.getInputs()) ||
110  containsSparseTensor(funcType.getResults()))
111  return true;
112  }
113  return false;
114  });
115 
116  bufferization::BufferizationState bufferizationState;
117 
118  if (failed(bufferization::bufferizeModuleOp(getOperation(), updatedOptions,
119  bufferizationState)))
120  return failure();
121 
123  return success();
124  }
125 
126  void runOnOperation() override {
127  // Overrides the default emit strategy using user-provided value.
128  this->sparsificationOptions.sparseEmitStrategy = sparseEmitStrategy;
129 
130  // Overrides the default parallelization strategy using user-provided value.
131  this->sparsificationOptions.parallelizationStrategy = parallelization;
132 
133  // Run enabling transformations.
134  {
135  OpPassManager pm("builtin.module");
137  pm.addNestedPass<func::FuncOp>(
138  bufferization::createEmptyTensorToAllocTensorPass());
139  if (failed(runPipeline(pm, getOperation())))
140  return signalPassFailure();
141  }
142 
143  // Insert tensor copies. This step runs One-Shot Analysis (which analyzes
144  // SSA use-def chains of tensor IR) and decides where buffer copies are
145  // needed and where buffers can be written to in-place. These decisions are
146  // materialized in the IR in the form of `bufferization.alloc_tensor` ops.
147  //
148  // Note: All following steps in this pass must be careful not to modify the
149  // structure of the IR (i.e., tensor use-def chains), as that could
150  // invalidate the results of the analysis. From now on, only small and
151  // localized rewrites are allowed, such as replacing a tensor op with its
152  // memref equivalent.
153  bufferization::BufferizationState bufferizationState;
154 
156  getOperation(), bufferizationOptions, bufferizationState)))
157  return signalPassFailure();
158 
159  // Option `testAnalysisOnly` is a debug/testing flag. If set, the results of
160  // OneShotAnalysis are added to the IR via attributes. In that case, do not
161  // continue with the remaining pipeline.
162  if (bufferizationOptions.testAnalysisOnly)
163  return;
164 
165  // Bufferize all sparse ops. No further analysis is needed. All required
166  // buffer copies were already inserted by `insertTensorCopies` in the form
167  // of `bufferization.alloc_tensor` ops.
168  {
169  OpPassManager pm("builtin.module");
170  if (enableGPULibgen)
171  pm.addPass(createSparseGPUCodegenPass(0, enableRuntimeLibrary));
173  pm.addPass(createSparsificationPass(sparsificationOptions));
174  if (sparsificationOptions.sparseEmitStrategy ==
176  pm.addNestedPass<func::FuncOp>(createSparseSpaceCollapsePass());
178  }
179 
181  pm.addPass(createLowerSparseOpsToForeachPass(enableRuntimeLibrary,
182  /*enableConvert=*/true));
183  pm.addPass(
185  pm.addNestedPass<func::FuncOp>(createLowerForeachToSCFPass());
187  if (vectorLength > 0) {
189  vectorLength, enableVLAVectorization, enableSIMDIndex32));
190  }
191  if (enableRuntimeLibrary) {
193  } else {
194  pm.addPass(createSparseTensorCodegenPass(createSparseDeallocs,
195  enableBufferInitialization));
196  pm.addPass(createSparseBufferRewritePass(enableBufferInitialization));
197  }
198  if (failed(runPipeline(pm, getOperation())))
199  return signalPassFailure();
200  }
201 
202  // Bufferize all dense ops.
203  if (failed(runDenseBufferization()))
204  signalPassFailure();
205  }
206 
207 private:
208  bufferization::OneShotBufferizationOptions bufferizationOptions;
209  SparsificationOptions sparsificationOptions;
210  bool createSparseDeallocs;
211  bool enableRuntimeLibrary;
212  bool enableBufferInitialization;
213 };
214 
215 } // namespace sparse_tensor
216 } // namespace mlir
217 
220  using namespace mlir::bufferization;
222  options.bufferizeFunctionBoundaries = true;
223  options.setFunctionBoundaryTypeConversion(LayoutMapOption::IdentityLayoutMap);
224  options.unknownTypeConverterFn = [](TensorType tensorType,
225  Attribute memorySpace,
226  const BufferizationOptions &options) {
227  return getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace);
228  };
229  if (analysisOnly) {
230  options.testAnalysisOnly = true;
231  options.printConflicts = true;
232  }
233  // Since this mini-pipeline may be used in alternative pipelines (viz.
234  // different from the default "sparsifier" pipeline) where unknown ops
235  // are handled by alternative bufferization methods that are downstream
236  // of this mini-pipeline, we allow unknown ops by default (failure to
237  // bufferize is eventually apparent by failing to convert to LLVM IR).
238  options.allowUnknownOps = true;
239  return options;
240 }
241 
242 std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass() {
243  SparsificationOptions sparseOptions;
244  return std::make_unique<
246  getBufferizationOptionsForSparsification(/*analysisOnly=*/false),
247  sparseOptions,
248  /*createSparseDeallocs=*/false,
249  /*enableRuntimeLibrary=*/false,
250  /*enableBufferInitialization=*/false);
251 }
252 
254  const bufferization::OneShotBufferizationOptions &bufferizationOptions,
255  const SparsificationOptions &sparsificationOptions,
256  bool createSparseDeallocs, bool enableRuntimeLibrary,
257  bool enableBufferInitialization, unsigned vectorLength,
258  bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen,
259  SparseEmitStrategy emitStrategy,
260  SparseParallelizationStrategy parallelizationStrategy) {
261  return std::make_unique<
263  bufferizationOptions, sparsificationOptions, createSparseDeallocs,
264  enableRuntimeLibrary, enableBufferInitialization, vectorLength,
265  enableVLAVectorization, enableSIMDIndex32, enableGPULibgen, emitStrategy,
266  parallelizationStrategy);
267 }
static llvm::ManagedStatic< PassManagerOptions > options
Attributes are known-constant values of operations.
Definition: Attributes.h:25
This class represents a pass manager that runs passes on either a specific operation type,...
Definition: PassManager.h:46
void addPass(std::unique_ptr< Pass > pass)
Add the given pass to this pass manager.
Definition: Pass.cpp:367
void addNestedPass(std::unique_ptr< Pass > pass)
Add the given pass to a nested pass manager for the given operation kind OpT.
Definition: PassManager.h:115
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
operand_range getOperands()
Returns an iterator on the underlying Value's.
Definition: Operation.h:378
result_range getResults()
Definition: Operation.h:415
Tensor types represent multi-dimensional arrays, and have two variants: RankedTensorType and Unranked...
Definition: BuiltinTypes.h:55
This class provides an abstraction over the various different ranges of value types.
Definition: TypeRange.h:37
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
BufferizationState provides information about the state of the IR during the bufferization process.
void denyOperation()
Deny the given ops.
A pass that lowers tensor ops to memref ops, regardless of whether they are dense or sparse.
SparsificationAndBufferizationPass(const bufferization::OneShotBufferizationOptions &bufferizationOptions, const SparsificationOptions &sparsificationOptions, bool createSparseDeallocs, bool enableRuntimeLibrary, bool enableBufferInitialization, unsigned vl, bool vla, bool index32, bool gpu, SparseEmitStrategy emitStrategy, SparseParallelizationStrategy parallelizationStrategy)
SparsificationAndBufferizationPass(const bufferization::OneShotBufferizationOptions &bufferizationOptions, const SparsificationOptions &sparsificationOptions, bool createSparseDeallocs, bool enableRuntimeLibrary, bool enableBufferInitialization)
BaseMemRefType getMemRefTypeWithStaticIdentityLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with a static identity layout (i.e., no layout map).
llvm::LogicalResult bufferizeModuleOp(Operation *moduleOp, const OneShotBufferizationOptions &options, BufferizationState &state, BufferizationStatistics *statistics=nullptr)
Bufferize an ops nested ops that implement BufferizableOpInterface.
void removeBufferizationAttributesInModule(Operation *moduleOp)
Remove bufferization attributes on every FuncOp arguments in the SymbolTable op.
LogicalResult insertTensorCopies(Operation *op, const OneShotBufferizationOptions &options, const BufferizationState &bufferizationState, BufferizationStatistics *statistics=nullptr)
Resolve RaW and other conflicts by inserting bufferization.alloc_tensor ops.
static bool containsSparseTensor(TypeRange types)
Return true if one of the given types is a sparse tensor type.
SparseTensorEncodingAttr getSparseTensorEncoding(Type type)
Convenience method to get a sparse encoding attribute from a type.
Include the generated interface declarations.
std::unique_ptr< Pass > createSparseVectorizationPass()
std::unique_ptr< Pass > createLowerSparseOpsToForeachPass()
std::unique_ptr< Pass > createSparseTensorCodegenPass()
std::unique_ptr< Pass > createSparseGPUCodegenPass()
std::unique_ptr< Pass > createSparseSpaceCollapsePass()
std::unique_ptr< Pass > createLoopInvariantCodeMotionPass()
Creates a loop invariant code motion pass that hoists loop invariant instructions out of the loop.
bufferization::OneShotBufferizationOptions getBufferizationOptionsForSparsification(bool analysisOnly)
std::unique_ptr< Pass > createSparseReinterpretMapPass()
std::unique_ptr< Pass > createSparseTensorConversionPass()
std::unique_ptr< Pass > createSparseBufferRewritePass()
std::unique_ptr< Pass > createSparsificationAndBufferizationPass()
SparseParallelizationStrategy
Defines a parallelization strategy.
Definition: Passes.h:36
SparseEmitStrategy
Defines a scope for reinterpret map pass.
Definition: Passes.h:52
std::unique_ptr< Pass > createPreSparsificationRewritePass()
std::unique_ptr< Pass > createLowerForeachToSCFPass()
std::unique_ptr< Pass > createLowerSparseIterationToSCFPass()
std::unique_ptr< Pass > createStageSparseOperationsPass()
std::unique_ptr< Pass > createSparsificationPass()
Options for the Sparsification pass.
Definition: Passes.h:93
SparseEmitStrategy sparseEmitStrategy
Definition: Passes.h:107
SparseParallelizationStrategy parallelizationStrategy
Definition: Passes.h:106
Options for BufferizableOpInterface-based bufferization.
bool testAnalysisOnly
If set to true, does not modify the IR apart from adding attributes (for checking the results of the ...
OpFilter opFilter
A filter that specifies which ops should be bufferized and which ops should be ignored.
Options for analysis-enabled bufferization.