MLIR  18.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:
64  const bufferization::OneShotBufferizationOptions &bufferizationOptions,
65  const SparsificationOptions &sparsificationOptions,
66  bool createSparseDeallocs, bool enableRuntimeLibrary,
67  bool enableBufferInitialization, unsigned vectorLength,
68  bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen)
69  : bufferizationOptions(bufferizationOptions),
70  sparsificationOptions(sparsificationOptions),
71  createSparseDeallocs(createSparseDeallocs),
72  enableRuntimeLibrary(enableRuntimeLibrary),
73  enableBufferInitialization(enableBufferInitialization),
74  vectorLength(vectorLength),
75  enableVLAVectorization(enableVLAVectorization),
76  enableSIMDIndex32(enableSIMDIndex32), enableGPULibgen(enableGPULibgen) {
77  }
78 
79  /// Bufferize all dense ops. This assumes that no further analysis is needed
80  /// and that all required buffer copies were already inserted by
81  /// `insertTensorCopies` in the form of `bufferization.alloc_tensor` ops.
84  bufferizationOptions;
85  // Skip all sparse ops.
86  updatedOptions.opFilter.denyOperation([&](Operation *op) {
89  return true;
90  if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
91  FunctionType funcType = funcOp.getFunctionType();
92  if (containsSparseTensor(funcType.getInputs()) ||
93  containsSparseTensor(funcType.getResults()))
94  return true;
95  }
96  return false;
97  });
98 
99  if (failed(bufferization::bufferizeModuleOp(cast<ModuleOp>(getOperation()),
100  updatedOptions)))
101  return failure();
102 
104  return success();
105  }
106 
107  void runOnOperation() override {
108  // Run enabling transformations.
109  {
110  OpPassManager pm("builtin.module");
112  pm.addNestedPass<func::FuncOp>(
114  if (failed(runPipeline(pm, getOperation())))
115  return signalPassFailure();
116  }
117 
118  // Insert tensor copies. This step runs One-Shot Analysis (which analyzes
119  // SSA use-def chains of tensor IR) and decides where buffer copies are
120  // needed and where buffers can be written to in-place. These decisions are
121  // materialized in the IR in the form of `bufferization.alloc_tensor` ops.
122  //
123  // Note: All following steps in this pass must be careful not to modify the
124  // structure of the IR (i.e., tensor use-def chains), as that could
125  // invalidate the results of the analysis. From now on, only small and
126  // localized rewrites are allowed, such as replacing a tensor op with its
127  // memref equivalent.
128  if (failed(bufferization::insertTensorCopies(getOperation(),
129  bufferizationOptions)))
130  return signalPassFailure();
131 
132  // Option `testAnalysisOnly` is a debug/testing flag. If set, the results of
133  // OneShotAnalysis are added to the IR via attributes. In that case, do not
134  // continue with the remaining pipeline.
135  if (bufferizationOptions.testAnalysisOnly)
136  return;
137 
138  // Bufferize all sparse ops. No further analysis is needed. All required
139  // buffer copies were already inserted by `insertTensorCopies` in the form
140  // of `bufferization.alloc_tensor` ops.
141  {
142  OpPassManager pm("builtin.module");
143  if (enableGPULibgen)
144  pm.addPass(createSparseGPUCodegenPass(0, enableRuntimeLibrary));
146  pm.addPass(createSparsificationPass(sparsificationOptions));
148  pm.addPass(createLowerSparseOpsToForeachPass(enableRuntimeLibrary,
149  /*enableConvert=*/true));
150  pm.addPass(
152  pm.addNestedPass<func::FuncOp>(createLowerForeachToSCFPass());
154  if (vectorLength > 0) {
156  vectorLength, enableVLAVectorization, enableSIMDIndex32));
157  }
158  if (enableRuntimeLibrary) {
160  } else {
161  pm.addPass(createSparseTensorCodegenPass(createSparseDeallocs,
162  enableBufferInitialization));
163  pm.addPass(createSparseBufferRewritePass(enableBufferInitialization));
164  }
165  if (failed(runPipeline(pm, getOperation())))
166  return signalPassFailure();
167  }
168 
169  // Bufferize all dense ops.
171  signalPassFailure();
172  }
173 
174 private:
175  bufferization::OneShotBufferizationOptions bufferizationOptions;
176  SparsificationOptions sparsificationOptions;
177  bool createSparseDeallocs;
178  bool enableRuntimeLibrary;
179  bool enableBufferInitialization;
180  unsigned vectorLength;
181  bool enableVLAVectorization;
182  bool enableSIMDIndex32;
183  bool enableGPULibgen;
184 };
185 
186 } // namespace sparse_tensor
187 } // namespace mlir
188 
191  using namespace mlir::bufferization;
193  options.bufferizeFunctionBoundaries = true;
194  options.setFunctionBoundaryTypeConversion(LayoutMapOption::IdentityLayoutMap);
195  options.unknownTypeConverterFn = [](Value value, Attribute memorySpace,
196  const BufferizationOptions &options) {
198  cast<TensorType>(value.getType()), memorySpace);
199  };
200  if (analysisOnly) {
201  options.testAnalysisOnly = true;
202  options.printConflicts = true;
203  }
204  return options;
205 }
206 
207 std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass() {
208  SparsificationOptions sparseOptions;
210  getBufferizationOptionsForSparsification(/*analysisOnly=*/false),
211  sparseOptions,
212  /*createSparseDeallocs=*/false,
213  /*enableRuntimeLibrary=*/false,
214  /*enableBufferInitialization=*/false,
215  /*vectorLength=*/0,
216  /*enableVLAVectorization=*/false,
217  /*enableSIMDIndex32=*/false,
218  /*enableGPULibgen=*/false);
219 }
220 
222  const bufferization::OneShotBufferizationOptions &bufferizationOptions,
223  const SparsificationOptions &sparsificationOptions,
224  bool createSparseDeallocs, bool enableRuntimeLibrary,
225  bool enableBufferInitialization, unsigned vectorLength,
226  bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen) {
227  return std::make_unique<
229  bufferizationOptions, sparsificationOptions, createSparseDeallocs,
230  enableRuntimeLibrary, enableBufferInitialization, vectorLength,
231  enableVLAVectorization, enableSIMDIndex32, enableGPULibgen);
232 }
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:48
void addPass(std::unique_ptr< Pass > pass)
Add the given pass to this pass manager.
Definition: Pass.cpp:346
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:117
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:373
result_range getResults()
Definition: Operation.h:410
This class provides an abstraction over the various different ranges of value types.
Definition: TypeRange.h:36
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Type getType() const
Return the type of this value.
Definition: Value.h:125
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 vectorLength, bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen)
BaseMemRefType getMemRefTypeWithStaticIdentityLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with a static identity layout (i.e., no layout map).
LogicalResult insertTensorCopies(Operation *op, const OneShotBufferizationOptions &options, BufferizationStatistics *statistics=nullptr)
Resolve RaW and other conflicts by inserting bufferization.alloc_tensor ops.
void removeBufferizationAttributesInModule(ModuleOp moduleOp)
Remove bufferization attributes on every FuncOp arguments in the ModuleOp.
std::unique_ptr< Pass > createEmptyTensorToAllocTensorPass()
Create a pass that rewrites tensor.empty to bufferization.alloc_tensor.
LogicalResult bufferizeModuleOp(ModuleOp moduleOp, const OneShotBufferizationOptions &options, BufferizationStatistics *statistics=nullptr)
Bufferize op and its nested ops that implement BufferizableOpInterface.
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()
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
std::unique_ptr< Pass > createLowerSparseOpsToForeachPass()
std::unique_ptr< Pass > createSparseTensorCodegenPass()
std::unique_ptr< Pass > createSparseGPUCodegenPass()
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()
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
std::unique_ptr< Pass > createPreSparsificationRewritePass()
std::unique_ptr< Pass > createLowerForeachToSCFPass()
std::unique_ptr< Pass > createStageSparseOperationsPass()
std::unique_ptr< Pass > createSparsificationPass()
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
Definition: LogicalResult.h:72
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
Options for the Sparsification pass.
Definition: Passes.h:76
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.