1 //===- SparseTensorPipelines.cpp - Pipelines for sparse tensor code -------===//
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 
9 #include "mlir/Dialect/SparseTensor/Pipelines/Passes.h"
10 
11 #include "mlir/Conversion/Passes.h"
12 #include "mlir/Dialect/Arithmetic/Transforms/Passes.h"
13 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
14 #include "mlir/Dialect/Func/IR/FuncOps.h"
15 #include "mlir/Dialect/Func/Transforms/Passes.h"
16 #include "mlir/Dialect/Linalg/Passes.h"
17 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
18 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
19 #include "mlir/Dialect/Tensor/Transforms/Passes.h"
20 #include "mlir/Dialect/Vector/Transforms/Passes.h"
21 #include "mlir/Pass/PassManager.h"
22 
23 using namespace mlir;
24 using namespace mlir::sparse_tensor;
25 
26 //===----------------------------------------------------------------------===//
27 // Pipeline implementation.
28 //===----------------------------------------------------------------------===//
29 
30 void mlir::sparse_tensor::buildSparseCompiler(
31     OpPassManager &pm, const SparseCompilerOptions &options) {
32   // TODO(wrengr): ensure the original `pm` is for ModuleOp
33   pm.addNestedPass<FuncOp>(createLinalgGeneralizationPass());
34   pm.addPass(createLinalgElementwiseOpFusionPass());
35   pm.addPass(createSparsificationPass(options.sparsificationOptions()));
36   pm.addPass(createSparseTensorConversionPass(
37       options.sparseTensorConversionOptions()));
38   pm.addNestedPass<FuncOp>(createLinalgBufferizePass());
39   pm.addNestedPass<FuncOp>(vector::createVectorBufferizePass());
40   pm.addNestedPass<FuncOp>(createConvertLinalgToLoopsPass());
41   pm.addNestedPass<FuncOp>(createConvertVectorToSCFPass());
42   pm.addNestedPass<FuncOp>(createConvertSCFToCFPass());
43   pm.addPass(func::createFuncBufferizePass());
44   pm.addPass(arith::createConstantBufferizePass());
45   pm.addNestedPass<FuncOp>(createTensorBufferizePass());
46   pm.addNestedPass<FuncOp>(
47       mlir::bufferization::createFinalizingBufferizePass());
48   pm.addPass(createLowerAffinePass());
49   pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions()));
50   pm.addPass(createMemRefToLLVMPass());
51   pm.addNestedPass<FuncOp>(createConvertMathToLLVMPass());
52   pm.addPass(createConvertFuncToLLVMPass());
53   pm.addPass(createReconcileUnrealizedCastsPass());
54 }
55 
56 //===----------------------------------------------------------------------===//
57 // Pipeline registration.
58 //===----------------------------------------------------------------------===//
59 
60 void mlir::sparse_tensor::registerSparseTensorPipelines() {
61   PassPipelineRegistration<SparseCompilerOptions>(
62       "sparse-compiler",
63       "The standard pipeline for taking sparsity-agnostic IR using the"
64       " sparse-tensor type, and lowering it to LLVM IR with concrete"
65       " representations and algorithms for sparse tensors.",
66       buildSparseCompiler);
67 }
68