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