150 lines · cpp
1//===- GPUToXeVMPipeline.cpp - Lowering pipeline to XeVM/LLVM -------------===//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// This file implements a pass for testing the lowering to XeVM as a generally10// usable sink pass. If XeGPU ops are used, it expects the MLIR code to have11// XeGPU ops already embedded in gpu code.12//13//===----------------------------------------------------------------------===//14 15#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"16#include "mlir/Conversion/MathToXeVM/MathToXeVM.h"17#include "mlir/Conversion/Passes.h"18#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"19#include "mlir/Conversion/VectorToSCF/VectorToSCF.h"20#include "mlir/Conversion/XeGPUToXeVM/XeGPUToXeVM.h"21#include "mlir/Conversion/XeVMToLLVM/XeVMToLLVM.h"22#include "mlir/Dialect/Func/IR/FuncOps.h"23#include "mlir/Dialect/GPU/IR/GPUDialect.h"24#include "mlir/Dialect/GPU/Pipelines/Passes.h"25#include "mlir/Dialect/GPU/Transforms/Passes.h"26#include "mlir/Dialect/LLVMIR/Transforms/RequestCWrappers.h"27#include "mlir/Dialect/MemRef/Transforms/Passes.h"28#include "mlir/Dialect/XeGPU/Transforms/Passes.h"29#include "mlir/Pass/PassManager.h"30#include "mlir/Pass/PassOptions.h"31#include "mlir/Target/LLVM/XeVM/Target.h"32#include "mlir/Transforms/Passes.h"33 34using namespace mlir;35 36namespace {37//===----------------------------------------------------------------------===//38// Pre-GPU common pipeline for both Host and GPU.39//===----------------------------------------------------------------------===//40void buildPreGPUCommonPassPipeline(41 OpPassManager &pm, const mlir::gpu::GPUToXeVMPipelineOptions &options) {42 // builtin.module scope passes.43 pm.addPass(createCSEPass());44 pm.addPass(createConvertVectorToSCFPass());45 {46 GpuXeVMAttachTargetOptions xevmTargetOptions;47 xevmTargetOptions.moduleMatcher = options.xevmModuleMatcher;48 xevmTargetOptions.triple = options.zebinTriple;49 xevmTargetOptions.chip = options.zebinChip;50 xevmTargetOptions.optLevel = options.optLevel;51 xevmTargetOptions.cmdOptions = options.cmdOptions;52 pm.addPass(createGpuXeVMAttachTarget(xevmTargetOptions));53 }54 pm.addPass(createLowerAffinePass());55 pm.addNestedPass<func::FuncOp>(createGpuAsyncRegionPass());56}57 58//===----------------------------------------------------------------------===//59// GPUModule-specific stuff.60//===----------------------------------------------------------------------===//61void buildGPUPassPipeline(OpPassManager &pm,62 const mlir::gpu::GPUToXeVMPipelineOptions &options) {63 if (options.xegpuOpLevel == "workgroup") {64 pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUWgToSgDistribute());65 pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());66 xegpu::XeGPUPropagateLayoutOptions layoutOptions;67 layoutOptions.layoutKind = "inst";68 pm.addNestedPass<gpu::GPUModuleOp>(69 xegpu::createXeGPUPropagateLayout(layoutOptions));70 pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUBlocking());71 pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());72 pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());73 }74 if (options.xegpuOpLevel == "subgroup" ||75 options.xegpuOpLevel == "workgroup") {76 xegpu::XeGPUPropagateLayoutOptions layoutOptions;77 layoutOptions.layoutKind = "lane";78 pm.addNestedPass<gpu::GPUModuleOp>(79 xegpu::createXeGPUPropagateLayout(layoutOptions));80 pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUSubgroupDistribute());81 pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());82 pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());83 pm.addNestedPass<gpu::GPUModuleOp>(createLoopInvariantCodeMotionPass());84 pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());85 pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUVectorLinearize());86 }87 pm.addNestedPass<gpu::GPUModuleOp>(createConvertMathToXeVM());88 pm.addNestedPass<gpu::GPUModuleOp>(createConvertXeGPUToXeVMPass());89 {90 ConvertGpuOpsToLLVMSPVOpsOptions gpuToLLVMSPVOptions;91 gpuToLLVMSPVOptions.use64bitIndex = options.use64bitIndex;92 pm.addNestedPass<gpu::GPUModuleOp>(93 createConvertGpuOpsToLLVMSPVOps(gpuToLLVMSPVOptions));94 }95 pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());96 pm.addNestedPass<gpu::GPUModuleOp>(createReconcileUnrealizedCastsPass());97}98 99//===----------------------------------------------------------------------===//100// Post-GPU pipeline for both Host and GPU.101//===----------------------------------------------------------------------===//102void buildPostGPUCommonPassPipeline(103 OpPassManager &pm, const mlir::gpu::GPUToXeVMPipelineOptions &options) {104 // builtin.module scope passes.105 pm.addPass(createSCFToControlFlowPass());106 pm.addPass(memref::createExpandStridedMetadataPass());107 {108 GpuToLLVMConversionPassOptions gpuToLLVMOptions;109 gpuToLLVMOptions.hostBarePtrCallConv = options.hostBarePtrCallConv;110 gpuToLLVMOptions.kernelBarePtrCallConv = options.kernelBarePtrCallConv;111 pm.addPass(createGpuToLLVMConversionPass(gpuToLLVMOptions));112 }113 pm.addPass(createLowerAffinePass());114 pm.addPass(createConvertVectorToLLVMPass());115 pm.addPass(createConvertToLLVMPass());116 pm.addPass(createReconcileUnrealizedCastsPass());117 pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());118 pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());119 // gpu-module-to-binary120 {121 GpuModuleToBinaryPassOptions gpuToModuleBinOptions;122 gpuToModuleBinOptions.compilationTarget = options.binaryFormat;123 gpuToModuleBinOptions.cmdOptions = options.cmdOptions;124 pm.addPass(createGpuModuleToBinaryPass(gpuToModuleBinOptions));125 }126}127} // namespace128 129void mlir::gpu::buildLowerToXeVMPassPipeline(130 OpPassManager &pm, const GPUToXeVMPipelineOptions &options) {131 // Pre-GPU common pipelines.132 buildPreGPUCommonPassPipeline(pm, options);133 134 // GPUModule-specific stuff.135 buildGPUPassPipeline(pm, options);136 137 // Post-GPU pipeline for both Host and GPU.138 buildPostGPUCommonPassPipeline(pm, options);139}140 141void mlir::gpu::registerGPUToXeVMPipeline() {142 PassPipelineRegistration<GPUToXeVMPipelineOptions>(143 "gpu-lower-to-xevm-pipeline",144 "The default GPU to XeVM lowering pipeline. It starts by lowering GPU "145 "code to the "146 "specified compilation target (default is fatbin) then lowers the host "147 "code.",148 buildLowerToXeVMPassPipeline);149}150