warpforth/lib/Conversion/Passes.cpp at 32eef876c555e3963c3c9aa1351328c25410fb1f · tetsuo-cpp/warpforth · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
//===- Passes.cpp - Conversion Pass Registration ----------------*- C++ -*-===//
//
// This file implements pass registration for conversion passes.
//
//===----------------------------------------------------------------------===//

#include "warpforth/Conversion/Passes.h"
#include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
#include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
#include "mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h"
#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/GPU/Transforms/Passes.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
#include "mlir/Transforms/Passes.h"
#include "warpforth/Conversion/ForthToGPU/ForthToGPU.h"
#include "warpforth/Conversion/ForthToMemRef/ForthToMemRef.h"

namespace mlir {
namespace warpforth {

void buildWarpForthPipeline(OpPassManager &pm) {
  // Stage 1: Lower Forth to MemRef (CF ops pass through as-is)
  pm.addPass(createConvertForthToMemRefPass());

  // Stage 2: Convert to GPU dialect (includes private address space annotation)
  pm.addPass(createConvertForthToGPUPass());

  // Stage 3: Normalize MemRefs for GPU
  pm.addPass(createCanonicalizerPass());

  // Stage 4: Attach NVVM target to GPU modules (sm_70 = Volta architecture)
  pm.addPass(createGpuNVVMAttachTarget());

  // Stage 5: Lower GPU to NVVM with bare pointers
  ConvertGpuOpsToNVVMOpsOptions gpuToNVVMOptions;
  gpuToNVVMOptions.useBarePtrCallConv = true;
  pm.addNestedPass<gpu::GPUModuleOp>(
      createConvertGpuOpsToNVVMOps(gpuToNVVMOptions));

  // Stage 6: Lower math ops to LLVM intrinsics inside GPU module
  pm.addNestedPass<gpu::GPUModuleOp>(createConvertMathToLLVMPass());

  // Stage 7: Lower NVVM to LLVM
  pm.addPass(createConvertNVVMToLLVMPass());

  // Stage 8: Reconcile type conversions
  pm.addPass(createReconcileUnrealizedCastsPass());

  // Stage 9: Compile GPU module to PTX binary
  GpuModuleToBinaryPassOptions binaryOptions;
  binaryOptions.compilationTarget = "isa"; // Output PTX assembly
  pm.addPass(createGpuModuleToBinaryPass(binaryOptions));
}

void registerConversionPasses() {
  registerPass([]() -> std::unique_ptr<Pass> {
    return createConvertForthToMemRefPass();
  });
  registerPass(
      []() -> std::unique_ptr<Pass> { return createConvertForthToGPUPass(); });

  // Register WarpForth pipeline
  PassPipelineRegistration<>("warpforth-pipeline",
                             "WarpForth compilation pipeline (Forth to PTX)",
                             buildWarpForthPipeline);
}

} // namespace warpforth
} // namespace mlir