Page MenuHomePhabricator

Please use GitHub pull requests for new patches. Avoid migrating existing patches. Phabricator shutdown timeline

Differential D147010 Diff 508837 mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp

Changeset View

Standalone View

View Options

mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp

Show First 20 Lines • Show All 50 Lines • ▼ Show 20 Lines
class SparsificationAndBufferizationPass class SparsificationAndBufferizationPass
: public PassWrapper<SparsificationAndBufferizationPass, : public PassWrapper<SparsificationAndBufferizationPass,
OperationPass<ModuleOp>> { OperationPass<ModuleOp>> {
public: public:
SparsificationAndBufferizationPass( SparsificationAndBufferizationPass(
const bufferization::OneShotBufferizationOptions &bufferizationOptions, const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions, const SparsificationOptions &sparsificationOptions,
const SparseTensorConversionOptions &sparseTensorConversionOptions, const SparseTensorConversionOptions &sparseTensorConversionOptions,
bool enableRuntimeLibrary, bool enableBufferInitialization, bool createSparseDeallocs, bool enableRuntimeLibrary,
unsigned vectorLength, bool enableVLAVectorization, bool enableBufferInitialization, unsigned vectorLength,
bool enableSIMDIndex32) bool enableVLAVectorization, bool enableSIMDIndex32)
: bufferizationOptions(bufferizationOptions), : bufferizationOptions(bufferizationOptions),
sparsificationOptions(sparsificationOptions), sparsificationOptions(sparsificationOptions),
sparseTensorConversionOptions(sparseTensorConversionOptions), sparseTensorConversionOptions(sparseTensorConversionOptions),
createSparseDeallocs(createSparseDeallocs),
enableRuntimeLibrary(enableRuntimeLibrary), enableRuntimeLibrary(enableRuntimeLibrary),
enableBufferInitialization(enableBufferInitialization), enableBufferInitialization(enableBufferInitialization),
vectorLength(vectorLength), vectorLength(vectorLength),
enableVLAVectorization(enableVLAVectorization), enableVLAVectorization(enableVLAVectorization),
enableSIMDIndex32(enableSIMDIndex32) {} enableSIMDIndex32(enableSIMDIndex32) {}
/// Bufferize all dense ops. This assumes that no further analysis is needed /// Bufferize all dense ops. This assumes that no further analysis is needed
/// and that all required buffer copies were already inserted by /// and that all required buffer copies were already inserted by
▲ Show 20 Lines • Show All 69 Lines • ▼ Show 20 Lines // of `bufferization.alloc_tensor` ops.
pm.addPass(mlir::createLoopInvariantCodeMotionPass()); pm.addPass(mlir::createLoopInvariantCodeMotionPass());
pm.addPass(createSparseVectorizationPass( pm.addPass(createSparseVectorizationPass(
vectorLength, enableVLAVectorization, enableSIMDIndex32)); vectorLength, enableVLAVectorization, enableSIMDIndex32));
} }
if (enableRuntimeLibrary) { if (enableRuntimeLibrary) {
pm.addPass( pm.addPass(
createSparseTensorConversionPass(sparseTensorConversionOptions)); createSparseTensorConversionPass(sparseTensorConversionOptions));
} else { } else {
pm.addPass(createSparseTensorCodegenPass(enableBufferInitialization)); pm.addPass(createSparseTensorCodegenPass(createSparseDeallocs,
enableBufferInitialization));
pm.addPass(createSparseBufferRewritePass(enableBufferInitialization)); pm.addPass(createSparseBufferRewritePass(enableBufferInitialization));
pm.addPass(createStorageSpecifierToLLVMPass()); pm.addPass(createStorageSpecifierToLLVMPass());
} }
if (failed(runPipeline(pm, getOperation()))) if (failed(runPipeline(pm, getOperation())))
return signalPassFailure(); return signalPassFailure();
} }
// Bufferize all dense ops. // Bufferize all dense ops.
if (failed(runDenseBufferization())) if (failed(runDenseBufferization()))
signalPassFailure(); signalPassFailure();
} }
private: private:
bufferization::OneShotBufferizationOptions bufferizationOptions; bufferization::OneShotBufferizationOptions bufferizationOptions;
SparsificationOptions sparsificationOptions; SparsificationOptions sparsificationOptions;
SparseTensorConversionOptions sparseTensorConversionOptions; SparseTensorConversionOptions sparseTensorConversionOptions;
bool createSparseDeallocs;
bool enableRuntimeLibrary; bool enableRuntimeLibrary;
bool enableBufferInitialization; bool enableBufferInitialization;
unsigned vectorLength; unsigned vectorLength;
bool enableVLAVectorization; bool enableVLAVectorization;
bool enableSIMDIndex32; bool enableSIMDIndex32;
}; };
} // namespace sparse_tensor } // namespace sparse_tensor
} // namespace mlir } // namespace mlir
std::unique_ptr<Pass> mlir::createSparsificationAndBufferizationPass( std::unique_ptr<Pass> mlir::createSparsificationAndBufferizationPass(
const bufferization::OneShotBufferizationOptions &bufferizationOptions, const bufferization::OneShotBufferizationOptions &bufferizationOptions,
const SparsificationOptions &sparsificationOptions, const SparsificationOptions &sparsificationOptions,
const SparseTensorConversionOptions &sparseTensorConversionOptions, const SparseTensorConversionOptions &sparseTensorConversionOptions,
bool enableRuntimeLibrary, bool enableBufferInitialization, bool createSparseDeallocs, bool enableRuntimeLibrary,
unsigned vectorLength, bool enableVLAVectorization, bool enableBufferInitialization, unsigned vectorLength,
bool enableSIMDIndex32) { bool enableVLAVectorization, bool enableSIMDIndex32) {
return std::make_unique< return std::make_unique<
mlir::sparse_tensor::SparsificationAndBufferizationPass>( mlir::sparse_tensor::SparsificationAndBufferizationPass>(
bufferizationOptions, sparsificationOptions, bufferizationOptions, sparsificationOptions,
sparseTensorConversionOptions, enableRuntimeLibrary, sparseTensorConversionOptions, createSparseDeallocs, enableRuntimeLibrary,
enableBufferInitialization, vectorLength, enableVLAVectorization, enableBufferInitialization, vectorLength, enableVLAVectorization,
enableSIMDIndex32); enableSIMDIndex32);
} }