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/GPUToNVVM/GPUToNVVMPass.h" 12 #include "mlir/Conversion/Passes.h" 13 #include "mlir/Dialect/Arith/Transforms/Passes.h" 14 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h" 15 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h" 16 #include "mlir/Dialect/Bufferization/Transforms/Passes.h" 17 #include "mlir/Dialect/Func/IR/FuncOps.h" 18 #include "mlir/Dialect/GPU/IR/GPUDialect.h" 19 #include "mlir/Dialect/GPU/Transforms/Passes.h" 20 #include "mlir/Dialect/LLVMIR/NVVMDialect.h" 21 #include "mlir/Dialect/Linalg/Passes.h" 22 #include "mlir/Dialect/MemRef/Transforms/Passes.h" 23 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h" 24 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h" 25 #include "mlir/Pass/PassManager.h" 26 #include "mlir/Transforms/Passes.h" 27 28 //===----------------------------------------------------------------------===// 29 // Pipeline implementation. 30 //===----------------------------------------------------------------------===// 31 32 void mlir::sparse_tensor::buildSparsifier(OpPassManager &pm, 33 const SparsifierOptions &options) { 34 // Rewrite named linalg ops into generic ops. 35 36 pm.addNestedPass<func::FuncOp>(createLinalgGeneralizationPass()); 37 38 // Sparsification and bufferization mini-pipeline. 39 pm.addPass(createSparsificationAndBufferizationPass( 40 getBufferizationOptionsForSparsification( 41 options.testBufferizationAnalysisOnly), 42 options.sparsificationOptions(), options.createSparseDeallocs, 43 options.enableRuntimeLibrary, options.enableBufferInitialization, 44 options.vectorLength, 45 /*enableVLAVectorization=*/options.armSVE, 46 /*enableSIMDIndex32=*/options.force32BitVectorIndices, 47 options.enableGPULibgen)); 48 49 // Bail-early for test setup. 50 if (options.testBufferizationAnalysisOnly) 51 return; 52 53 // Storage specifier lowering and bufferization wrap-up. 54 pm.addPass(createStorageSpecifierToLLVMPass()); 55 pm.addNestedPass<func::FuncOp>(createCanonicalizerPass()); 56 pm.addNestedPass<func::FuncOp>( 57 mlir::bufferization::createFinalizingBufferizePass()); 58 59 // GPU code generation. 60 const bool gpuCodegen = options.gpuTriple.hasValue(); 61 if (gpuCodegen) { 62 pm.addPass(createSparseGPUCodegenPass()); 63 pm.addNestedPass<gpu::GPUModuleOp>(createStripDebugInfoPass()); 64 pm.addNestedPass<gpu::GPUModuleOp>(createConvertSCFToCFPass()); 65 pm.addNestedPass<gpu::GPUModuleOp>(createConvertGpuOpsToNVVMOps()); 66 } 67 68 // TODO(springerm): Add sparse support to the BufferDeallocation pass and add 69 // it to this pipeline. 70 pm.addNestedPass<func::FuncOp>(createConvertLinalgToLoopsPass()); 71 pm.addNestedPass<func::FuncOp>(createConvertVectorToSCFPass()); 72 pm.addNestedPass<func::FuncOp>(memref::createExpandReallocPass()); 73 pm.addNestedPass<func::FuncOp>(createConvertSCFToCFPass()); 74 pm.addPass(memref::createExpandStridedMetadataPass()); 75 pm.addPass(createLowerAffinePass()); 76 pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions())); 77 pm.addPass(createFinalizeMemRefToLLVMConversionPass()); 78 pm.addNestedPass<func::FuncOp>(createConvertComplexToStandardPass()); 79 pm.addNestedPass<func::FuncOp>(arith::createArithExpandOpsPass()); 80 pm.addNestedPass<func::FuncOp>(createConvertMathToLLVMPass()); 81 pm.addPass(createConvertMathToLibmPass()); 82 pm.addPass(createConvertComplexToLibmPass()); 83 84 // Repeat convert-vector-to-llvm. 85 pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions())); 86 87 pm.addPass(createConvertComplexToLLVMPass()); 88 pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions())); 89 pm.addPass(createConvertFuncToLLVMPass()); 90 91 // Finalize GPU code generation. 92 if (gpuCodegen) { 93 GpuNVVMAttachTargetOptions nvvmTargetOptions; 94 nvvmTargetOptions.triple = options.gpuTriple; 95 nvvmTargetOptions.chip = options.gpuChip; 96 nvvmTargetOptions.features = options.gpuFeatures; 97 pm.addPass(createGpuNVVMAttachTarget(nvvmTargetOptions)); 98 pm.addPass(createGpuToLLVMConversionPass()); 99 GpuModuleToBinaryPassOptions gpuModuleToBinaryPassOptions; 100 gpuModuleToBinaryPassOptions.compilationTarget = options.gpuFormat; 101 pm.addPass(createGpuModuleToBinaryPass(gpuModuleToBinaryPassOptions)); 102 } 103 104 pm.addPass(createReconcileUnrealizedCastsPass()); 105 } 106 107 //===----------------------------------------------------------------------===// 108 // Pipeline registration. 109 //===----------------------------------------------------------------------===// 110 111 void mlir::sparse_tensor::registerSparseTensorPipelines() { 112 PassPipelineRegistration<SparsifierOptions>( 113 "sparsifier", 114 "The standard pipeline for taking sparsity-agnostic IR using the" 115 " sparse-tensor type, and lowering it to LLVM IR with concrete" 116 " representations and algorithms for sparse tensors.", 117 buildSparsifier); 118 } 119