MLIR  17.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 
22 #include "mlir/Pass/PassManager.h"
23 #include "mlir/Transforms/Passes.h"
24 
25 using namespace mlir;
26 using namespace mlir::func;
27 
28 namespace mlir {
29 namespace sparse_tensor {
30 
31 /// Return `true` if one of the given types is a sparse tensor type.
32 static bool containsSparseTensor(TypeRange types) {
33  for (Type t : types)
35  return true;
36  return false;
37 }
38 
39 /// A pass that lowers tensor ops to memref ops, regardless of whether they are
40 /// dense or sparse.
41 ///
42 /// One-Shot Analysis is used to detect RaW conflicts and to insert buffer
43 /// copies of the tensor level (`insertTensorCopies`). Afterwards, the lowering
44 /// of tensor ops to memref ops follows a different code path depending on
45 /// whether the op is sparse or dense:
46 ///
47 /// * Sparse tensor ops are lowered through Sparsification and follow-up pass
48 /// that lowers sparse_tensor dialect ops.
49 /// * Dense tensor ops are lowered through BufferizableOpInterface
50 /// implementations.
52  : public PassWrapper<SparsificationAndBufferizationPass,
53  OperationPass<ModuleOp>> {
54 public:
56  const bufferization::OneShotBufferizationOptions &bufferizationOptions,
57  const SparsificationOptions &sparsificationOptions,
58  const SparseTensorConversionOptions &sparseTensorConversionOptions,
59  bool createSparseDeallocs, bool enableRuntimeLibrary,
60  bool enableBufferInitialization, unsigned vectorLength,
61  bool enableVLAVectorization, bool enableSIMDIndex32)
62  : bufferizationOptions(bufferizationOptions),
63  sparsificationOptions(sparsificationOptions),
64  sparseTensorConversionOptions(sparseTensorConversionOptions),
65  createSparseDeallocs(createSparseDeallocs),
66  enableRuntimeLibrary(enableRuntimeLibrary),
67  enableBufferInitialization(enableBufferInitialization),
68  vectorLength(vectorLength),
69  enableVLAVectorization(enableVLAVectorization),
70  enableSIMDIndex32(enableSIMDIndex32) {}
71 
72  /// Bufferize all dense ops. This assumes that no further analysis is needed
73  /// and that all required buffer copies were already inserted by
74  /// `insertTensorCopies` in the form of `bufferization.alloc_tensor` ops.
76  bufferization::OpFilter denseOpFilter;
77  denseOpFilter.allowOperation([&](Operation *op) {
80  return false;
81  if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
82  FunctionType funcType = funcOp.getFunctionType();
83  if (containsSparseTensor(funcType.getInputs()) ||
84  containsSparseTensor(funcType.getResults()))
85  return false;
86  }
87  return true;
88  });
89 
90  if (failed(bufferization::bufferizeOp(getOperation(), bufferizationOptions,
91  /*copyBeforeWrite=*/false,
92  &denseOpFilter)))
93  return failure();
94 
96  return success();
97  }
98 
99  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
100  registry.insert<bufferization::BufferizationDialect>();
101  registry.insert<LLVM::LLVMDialect>();
102  }
103 
104  void runOnOperation() override {
105  {
106  // Run enabling transformations.
107  OpPassManager pm("builtin.module");
109  pm.addNestedPass<func::FuncOp>(
111  if (failed(runPipeline(pm, getOperation())))
112  return signalPassFailure();
113  }
114 
115  // Insert tensor copies. This step runs One-Shot Analysis (which analyzes
116  // SSA use-def chains of tensor IR) and decides where buffer copies are
117  // needed and where buffers can be written to in-place. These decisions are
118  // materialized in the IR in the form of `bufferization.alloc_tensor` ops.
119  //
120  // Note: All following steps in this pass must be careful not to modify the
121  // structure of the IR (i.e., tensor use-def chains), as that could
122  // invalidate the results of the analysis. From now on, only small and
123  // localized rewrites are allowed, such as replacing a tensor op with its
124  // memref equivalent.
125  if (failed(bufferization::insertTensorCopies(getOperation(),
126  bufferizationOptions)))
127  return signalPassFailure();
128 
129  // `testAnalysisOnly` is a debug/testing flag. If set, the results of
130  // OneShotAnalysis are added to the IR via attributes. In that case, do not
131  // continue with the remaining pipeline.
132  if (bufferizationOptions.testAnalysisOnly)
133  return;
134 
135  // Bufferize all sparse ops. No further analysis is needed. All required
136  // buffer copies were already inserted by `insertTensorCopies` in the form
137  // of `bufferization.alloc_tensor` ops.
138  {
139  OpPassManager pm("builtin.module");
140  pm.addPass(createSparsificationPass(sparsificationOptions));
141  pm.addPass(createPostSparsificationRewritePass(enableRuntimeLibrary));
142  if (vectorLength > 0) {
145  vectorLength, enableVLAVectorization, enableSIMDIndex32));
146  }
147  if (enableRuntimeLibrary) {
148  pm.addPass(
149  createSparseTensorConversionPass(sparseTensorConversionOptions));
150  } else {
151  pm.addPass(createSparseTensorCodegenPass(createSparseDeallocs,
152  enableBufferInitialization));
153  pm.addPass(createSparseBufferRewritePass(enableBufferInitialization));
155  }
156  if (failed(runPipeline(pm, getOperation())))
157  return signalPassFailure();
158  }
159 
160  // Bufferize all dense ops.
161  if (failed(runDenseBufferization()))
162  signalPassFailure();
163  }
164 
165 private:
166  bufferization::OneShotBufferizationOptions bufferizationOptions;
167  SparsificationOptions sparsificationOptions;
168  SparseTensorConversionOptions sparseTensorConversionOptions;
169  bool createSparseDeallocs;
170  bool enableRuntimeLibrary;
171  bool enableBufferInitialization;
172  unsigned vectorLength;
173  bool enableVLAVectorization;
174  bool enableSIMDIndex32;
175 };
176 
177 } // namespace sparse_tensor
178 } // namespace mlir
179 
181  const bufferization::OneShotBufferizationOptions &bufferizationOptions,
182  const SparsificationOptions &sparsificationOptions,
183  const SparseTensorConversionOptions &sparseTensorConversionOptions,
184  bool createSparseDeallocs, bool enableRuntimeLibrary,
185  bool enableBufferInitialization, unsigned vectorLength,
186  bool enableVLAVectorization, bool enableSIMDIndex32) {
187  return std::make_unique<
189  bufferizationOptions, sparsificationOptions,
190  sparseTensorConversionOptions, createSparseDeallocs, enableRuntimeLibrary,
191  enableBufferInitialization, vectorLength, enableVLAVectorization,
192  enableSIMDIndex32);
193 }
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
This class represents a pass manager that runs passes on either a specific operation type,...
Definition: PassManager.h:52
void addPass(std::unique_ptr< Pass > pass)
Add the given pass to this pass manager.
Definition: Pass.cpp:345
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:121
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:75
operand_range getOperands()
Returns an iterator on the underlying Value's.
Definition: Operation.h:357
result_range getResults()
Definition: Operation.h:394
This class provides a CRTP wrapper around a base pass class to define several necessary utility metho...
Definition: Pass.h:441
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
void allowOperation()
Allow the given ops.
A pass that lowers tensor ops to memref ops, regardless of whether they are dense or sparse.
void runOnOperation() override
The polymorphic API that runs the pass over the currently held operation.
void getDependentDialects(::mlir::DialectRegistry &registry) const override
SparsificationAndBufferizationPass(const bufferization::OneShotBufferizationOptions &bufferizationOptions, const SparsificationOptions &sparsificationOptions, const SparseTensorConversionOptions &sparseTensorConversionOptions, bool createSparseDeallocs, bool enableRuntimeLibrary, bool enableBufferInitialization, unsigned vectorLength, bool enableVLAVectorization, bool enableSIMDIndex32)
LogicalResult bufferizeOp(Operation *op, const BufferizationOptions &options, bool copyBeforeWrite=true, const OpFilter *opFilter=nullptr, BufferizationStatistics *statistics=nullptr)
Bufferize op and its nested ops that implement BufferizableOpInterface.
Definition: Bufferize.cpp:427
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.
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.
This header declares functions that assit transformations in the MemRef dialect.
std::unique_ptr< Pass > createSparseVectorizationPass()
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
std::unique_ptr< Pass > createSparseTensorCodegenPass()
std::unique_ptr< Pass > createLoopInvariantCodeMotionPass()
Creates a loop invariant code motion pass that hoists loop invariant instructions out of the loop.
std::unique_ptr< Pass > createSparsificationAndBufferizationPass(const bufferization::OneShotBufferizationOptions &bufferizationOptions, const SparsificationOptions &sparsificationOptions, const SparseTensorConversionOptions &sparseTensorConversionOptions, bool createSparseDeallocs, bool enableRuntimeLibrary, bool enableBufferInitialization, unsigned vectorLength, bool enableVLAVectorization, bool enableSIMDIndex32)
std::unique_ptr< Pass > createSparseTensorConversionPass()
std::unique_ptr< Pass > createSparseBufferRewritePass()
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
std::unique_ptr< Pass > createStorageSpecifierToLLVMPass()
std::unique_ptr< Pass > createPreSparsificationRewritePass()
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
std::unique_ptr< Pass > createPostSparsificationRewritePass()
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
SparseTensorConversion options.
Definition: Passes.h:103
Options for the Sparsification pass.
Definition: Passes.h:51
Options for analysis-enabled bufferization.