feat(warpforthc): add --emit to stop after mlir, llvm, or ptx

include/warpforth/Conversion/Passes.h

@@ -52,6 +52,11 @@ void buildWarpForthPipeline(OpPassManager &pm);
 void buildWarpForthPipeline(OpPassManager &pm,
                             const WarpForthPipelineOptions &options);
 
+/// Build the lowering stages of the WarpForth pipeline: Forth to a gpu.module
+/// whose body is in the LLVM/NVVM dialect, stopping before binary emission.
+void buildWarpForthLoweringPipeline(OpPassManager &pm,
+                                    const WarpForthPipelineOptions &options);
+
 } // namespace warpforth
 } // namespace mlir
 

lib/Conversion/Passes.cpp

@@ -87,8 +87,8 @@ static StringRef getBinaryFormat(gpu::CompilationTarget target) {
   llvm_unreachable("unknown gpu::CompilationTarget");
 }
 
-void buildWarpForthPipeline(OpPassManager &pm,
-                            const WarpForthPipelineOptions &options) {
+void buildWarpForthLoweringPipeline(OpPassManager &pm,
+                                    const WarpForthPipelineOptions &options) {
   // Stage 1: Lower Forth to MemRef (CF ops pass through as-is)
   pm.addPass(createConvertForthToMemRefPass());
 
@@ -132,6 +132,11 @@ void buildWarpForthPipeline(OpPassManager &pm,
 
   // Stage 9: Reconcile type conversions
   pm.addPass(createReconcileUnrealizedCastsPass());
+}
+
+void buildWarpForthPipeline(OpPassManager &pm,
+                            const WarpForthPipelineOptions &options) {
+  buildWarpForthLoweringPipeline(pm, options);
 
   // Stage 10: Compile GPU module to PTX binary
   GpuModuleToBinaryPassOptions binaryOptions;

test/Pipeline/warpforthc-emit.forth

@@ -0,0 +1,16 @@
+\ RUN: %warpforthc --emit=mlir %s | %FileCheck %s --check-prefix=MLIR
+\ RUN: %warpforthc --emit=llvm %s | %FileCheck %s --check-prefix=LLVM
+\ RUN: %warpforthc --emit=ptx %s | %FileCheck %s --check-prefix=PTX
+\ RUN: %not %warpforthc --emit=cubin %s 2>&1 | %FileCheck %s --check-prefix=BAD
+
+\! kernel main
+42
+
+\ MLIR: func.func @main
+\ MLIR: forth.constant
+\ MLIR-NOT: gpu.module
+\ LLVM: gpu.module @warpforth_module
+\ LLVM: #nvvm.target
+\ LLVM-NOT: gpu.binary
+\ PTX: .version
+\ BAD: Cannot find option named 'cubin'!

tools/warpforthc/warpforthc.cpp

@@ -59,6 +59,48 @@ static llvm::cl::opt<std::string>
     libdevicePath("libdevice", llvm::cl::desc("Path to the libdevice bitcode"),
                   llvm::cl::value_desc("path"), llvm::cl::init(""));
 
+enum class EmitKind { Mlir, Llvm, Ptx };
+
+// Not cl::Prefix: -O is a prefix option, and a prefix match on --emit would
+// also swallow LLVM's --emit-* flags.
+static llvm::cl::opt<EmitKind> emitKind(
+    "emit", llvm::cl::desc("Artifact to write instead of running every stage"),
+    llvm::cl::values(
+        clEnumValN(EmitKind::Mlir, "mlir", "Parsed Forth module as MLIR text"),
+        clEnumValN(EmitKind::Llvm, "llvm",
+                   "Lowered module in the LLVM dialect, as MLIR text"),
+        clEnumValN(EmitKind::Ptx, "ptx", "PTX assembly (default)")),
+    llvm::cl::init(EmitKind::Ptx));
+
+static void
+populatePipelineForEmit(PassManager &pm,
+                        const warpforth::WarpForthPipelineOptions &options,
+                        EmitKind kind) {
+  switch (kind) {
+  case EmitKind::Mlir:
+    return;
+  case EmitKind::Llvm:
+    warpforth::buildWarpForthLoweringPipeline(pm, options);
+    return;
+  case EmitKind::Ptx:
+    warpforth::buildWarpForthPipeline(pm, options);
+    return;
+  }
+}
+
+static LogicalResult writeEmitArtifact(ModuleOp module, EmitKind kind,
+                                       llvm::raw_ostream &os) {
+  switch (kind) {
+  case EmitKind::Mlir:
+  case EmitKind::Llvm:
+    module.print(os);
+    return success();
+  case EmitKind::Ptx:
+    return warpforth::extractPTXFromModule(module, os);
+  }
+  llvm_unreachable("unhandled EmitKind");
+}
+
 int main(int argc, char **argv) {
   llvm::InitLLVM y(argc, argv);
   warpforth::registerConversionPasses();
@@ -113,7 +155,7 @@ int main(int argc, char **argv) {
   if (failed(pipelineOptions.parseFromString(overrides)))
     return 1;
   pipelineOptions.chip = targetArchitecture.getValue();
-  warpforth::buildWarpForthPipeline(pm, pipelineOptions);
+  populatePipelineForEmit(pm, pipelineOptions, emitKind.getValue());
   if (failed(pm.run(*module)))
     return 1;
 
@@ -123,7 +165,7 @@ int main(int argc, char **argv) {
     return 1;
   }
 
-  if (failed(warpforth::extractPTXFromModule(*module, outputFile->os())))
+  if (failed(writeEmitArtifact(*module, emitKind.getValue(), outputFile->os())))
     return 1;
 
   outputFile->keep();