mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp

//===- GPUToXeVMPipeline.cpp - Lowering pipeline to XeVM/LLVM -------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements a pass for testing the lowering to XeVM as a generally
// usable sink pass. If XeGPU ops are used, it expects the MLIR code to have
// XeGPU ops already embedded in gpu code.
//
//===----------------------------------------------------------------------===//

#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
#include "mlir/Conversion/MathToXeVM/MathToXeVM.h"
#include "mlir/Conversion/Passes.h"
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
#include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
#include "mlir/Conversion/XeGPUToXeVM/XeGPUToXeVM.h"
#include "mlir/Conversion/XeVMToLLVM/XeVMToLLVM.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/GPU/Pipelines/Passes.h"
#include "mlir/Dialect/GPU/Transforms/Passes.h"
#include "mlir/Dialect/LLVMIR/Transforms/RequestCWrappers.h"
#include "mlir/Dialect/Math/Transforms/Passes.h"
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Pass/PassOptions.h"
#include "mlir/Target/LLVM/XeVM/Target.h"
#include "mlir/Transforms/Passes.h"

using namespace mlir;

namespace {
//===----------------------------------------------------------------------===//
// Pre-GPU common pipeline for both Host and GPU.
//===----------------------------------------------------------------------===//
void buildPreGPUCommonPassPipeline(
    OpPassManager &pm, const mlir::gpu::GPUToXeVMPipelineOptions &options) {
  // builtin.module scope passes.
  pm.addPass(createCSEPass());
  {
    GpuXeVMAttachTargetOptions xevmTargetOptions;
    xevmTargetOptions.moduleMatcher = options.xevmModuleMatcher;
    xevmTargetOptions.triple = options.zebinTriple;
    xevmTargetOptions.chip = options.zebinChip;
    xevmTargetOptions.optLevel = options.optLevel;
    xevmTargetOptions.cmdOptions = options.cmdOptions;
    pm.addPass(createGpuXeVMAttachTarget(xevmTargetOptions));
  }
  pm.addPass(createLowerAffinePass());
  pm.addNestedPass<func::FuncOp>(createGpuAsyncRegionPass());
}

//===----------------------------------------------------------------------===//
// GPUModule-specific stuff.
//===----------------------------------------------------------------------===//
void buildGPUPassPipeline(OpPassManager &pm,
                          const mlir::gpu::GPUToXeVMPipelineOptions &options) {
  xegpu::XeGPUPropagateLayoutOptions laneLayoutOptions;
  laneLayoutOptions.indexBitWidth = options.use64bitIndex ? 64 : 32;
  laneLayoutOptions.layoutKind = "lane";
  pm.addNestedPass<ModuleOp>(createCSEPass());
  if (options.enableVectorToXeGPU)
    pm.addNestedPass<gpu::GPUModuleOp>(createConvertVectorToXeGPU());
  if (options.xegpuOpLevel == "workgroup") {
    xegpu::XeGPUPropagateLayoutOptions sgLayoutOptions;
    sgLayoutOptions.layoutKind = "subgroup";
    pm.addNestedPass<gpu::GPUModuleOp>(
        xegpu::createXeGPUPropagateLayout(sgLayoutOptions));
    pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUWgToSgDistribute());
    pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
    pm.addNestedPass<gpu::GPUModuleOp>(createLowerAffinePass());
    pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
    xegpu::XeGPUPropagateLayoutOptions instDataOptions;
    instDataOptions.layoutKind = "inst";
    pm.addNestedPass<gpu::GPUModuleOp>(
        xegpu::createXeGPUPropagateLayout(instDataOptions));
    pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUBlocking());
    pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
  }
  if (options.xegpuOpLevel == "subgroup" ||
      options.xegpuOpLevel == "workgroup") {
    pm.addNestedPass<gpu::GPUModuleOp>(
        xegpu::createXeGPUPropagateLayout(laneLayoutOptions));
    pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUPeepHoleOptimizer());
    pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
    pm.addNestedPass<gpu::GPUModuleOp>(
        xegpu::createXeGPUPropagateLayout(laneLayoutOptions));
    pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUSgToLaneDistribute());
    pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
    pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
    pm.addNestedPass<gpu::GPUModuleOp>(createLoopInvariantCodeMotionPass());
    pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
    pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUVectorLinearize());
    pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
    pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
  }
  // Break down high-level micro-scaling (MX) ops (arith.scaling_extf and
  // arith.scaling_truncf) into standard arith ops (extf/truncf + mulf), and
  // expand extf/truncf on f8E8M0FNU into integer bit manipulation. This runs
  // before the XeVM/LLVM conversions. The f4E2M1FN expansion patterns are
  // intentionally left disabled: f4E2M1FN extf/truncf are lowered by the XeVM
  // conversions (xevm.extf), whereas f8E8M0FNU is not supported there and so
  // must be expanded here.
  {
    arith::ArithExpandOpsPassOptions arithExpandOptions;
    arithExpandOptions.includeF8E8M0 = true;
    pm.addNestedPass<gpu::GPUModuleOp>(
        arith::createArithExpandOpsPass(arithExpandOptions));
  }
  pm.addNestedPass<gpu::GPUModuleOp>(createConvertMathToXeVM());
  ConvertXeGPUToXeVMPassOptions xegpuToXeVMOptions;
  xegpuToXeVMOptions.use64bitIndex = options.use64bitIndex;
  pm.addNestedPass<gpu::GPUModuleOp>(
      createConvertXeGPUToXeVMPass(xegpuToXeVMOptions));
  {
    ConvertGpuOpsToLLVMSPVOpsOptions gpuToLLVMSPVOptions;
    gpuToLLVMSPVOptions.use64bitIndex = options.use64bitIndex;
    pm.addNestedPass<gpu::GPUModuleOp>(
        createConvertGpuOpsToLLVMSPVOps(gpuToLLVMSPVOptions));
  }
  // Legalize math/arith ops on floating-point types that the XeVM target
  // cannot handle natively (e.g. bf16) by wrapping them with extf/truncf
  // around a supported type (defaulting to f32).
  {
    math::MathExtendToSupportedTypesOptions mathExtendOptions;
    mathExtendOptions.extraTypeStrs.assign(options.mathExtendExtraTypes.begin(),
                                           options.mathExtendExtraTypes.end());
    mathExtendOptions.targetTypeStr = options.supportedTargetTypes;
    pm.addNestedPass<gpu::GPUModuleOp>(
        math::createMathExtendToSupportedTypes(mathExtendOptions));
  }
  {
    arith::ArithEmulateUnsupportedFloatsOptions arithEmulateOptions;
    arithEmulateOptions.sourceTypeStrs.assign(
        options.unsupportedSourceTypes.begin(),
        options.unsupportedSourceTypes.end());
    arithEmulateOptions.targetTypeStr = options.supportedTargetTypes;
    pm.addNestedPass<gpu::GPUModuleOp>(
        arith::createArithEmulateUnsupportedFloats(arithEmulateOptions));
  }
  pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
  pm.addNestedPass<gpu::GPUModuleOp>(createReconcileUnrealizedCastsPass());
}

//===----------------------------------------------------------------------===//
// Post-GPU pipeline for both Host and GPU.
//===----------------------------------------------------------------------===//
void buildPostGPUCommonPassPipeline(
    OpPassManager &pm, const mlir::gpu::GPUToXeVMPipelineOptions &options) {
  // builtin.module scope passes.
  pm.addPass(createConvertVectorToSCFPass());
  pm.addPass(createSCFToControlFlowPass());
  pm.addPass(memref::createExpandStridedMetadataPass());
  {
    GpuToLLVMConversionPassOptions gpuToLLVMOptions;
    gpuToLLVMOptions.hostBarePtrCallConv = options.hostBarePtrCallConv;
    gpuToLLVMOptions.kernelBarePtrCallConv = options.kernelBarePtrCallConv;
    pm.addPass(createGpuToLLVMConversionPass(gpuToLLVMOptions));
  }
  pm.addPass(createLowerAffinePass());
  pm.addPass(createConvertVectorToLLVMPass());
  pm.addPass(createConvertToLLVMPass());
  pm.addPass(createReconcileUnrealizedCastsPass());
  pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
  pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
  // XeVM-to-LLVM must be the last pass before gpu-module-to-binary.
  pm.addNestedPass<gpu::GPUModuleOp>(createConvertXeVMToLLVMPass());
  // gpu-module-to-binary
  {
    GpuModuleToBinaryPassOptions gpuToModuleBinOptions;
    gpuToModuleBinOptions.compilationTarget = options.binaryFormat;
    gpuToModuleBinOptions.cmdOptions = options.cmdOptions;
    pm.addPass(createGpuModuleToBinaryPass(gpuToModuleBinOptions));
  }
}
} // namespace

void mlir::gpu::buildLowerToXeVMPassPipeline(
    OpPassManager &pm, const GPUToXeVMPipelineOptions &options) {
  // Pre-GPU common pipelines.
  buildPreGPUCommonPassPipeline(pm, options);

  // GPUModule-specific stuff.
  buildGPUPassPipeline(pm, options);

  // Post-GPU pipeline for both Host and GPU.
  buildPostGPUCommonPassPipeline(pm, options);
}

void mlir::gpu::registerGPUToXeVMPipeline() {
  PassPipelineRegistration<GPUToXeVMPipelineOptions>(
      "gpu-lower-to-xevm-pipeline",
      "The default GPU to XeVM lowering pipeline. It starts by lowering GPU "
      "code to the "
      "specified compilation target (default is fatbin) then lowers the host "
      "code.",
      buildLowerToXeVMPassPipeline);
}