From 6a9685f580ac37270808ba37a2ddc8f20030b616 Mon Sep 17 00:00:00 2001
From: Alex Cameron <asc@tetsuo.sh>
Date: Thu, 19 Feb 2026 20:37:29 +0900
Subject: [PATCH 1/2] feat: replace PARAM with \! kernel header syntax

---
 CLAUDE.md                                     |   2 +-
 gpu_test/conftest.py                          |  70 ++++-
 gpu_test/test_kernels.py                      |  98 +++++--
 lib/Conversion/ForthToGPU/ForthToGPU.cpp      |   2 +-
 .../ForthToMemRef/ForthToMemRef.cpp           |  21 +-
 lib/Translation/ForthToMLIR/ForthToMLIR.cpp   | 243 +++++++++++++++---
 lib/Translation/ForthToMLIR/ForthToMLIR.h     |  25 +-
 .../Conversion/ForthToGPU/basic-gpu-wrap.mlir |   2 +-
 .../ForthToGPU/intrinsic-conversion.mlir      |   2 +-
 .../ForthToGPU/private-functions.mlir         |   2 +-
 test/Pipeline/begin-until.forth               |   3 +-
 test/Pipeline/begin-while-repeat.forth        |   3 +-
 test/Pipeline/control-flow.forth              |   3 +-
 test/Pipeline/do-loop.forth                   |   3 +-
 test/Pipeline/exit.forth                      |   3 +-
 test/Pipeline/full-pipeline.forth             |   3 +-
 test/Pipeline/interleaved-control-flow.forth  |   3 +-
 test/Pipeline/leave.forth                     |   3 +-
 test/Pipeline/matmul-naive.forth              |   7 +-
 test/Pipeline/multi-param.forth               |   5 +-
 test/Pipeline/nested-control-flow.forth       |   3 +-
 test/Pipeline/plus-loop-negative.forth        |   3 +-
 test/Pipeline/plus-loop.forth                 |   3 +-
 test/Pipeline/unloop-exit.forth               |   3 +-
 test/Translation/Forth/arithmetic-ops.forth   |   1 +
 test/Translation/Forth/basic-literals.forth   |   1 +
 test/Translation/Forth/begin-until.forth      |   1 +
 .../Forth/begin-while-repeat.forth            |   1 +
 test/Translation/Forth/bitwise-ops.forth      |   1 +
 test/Translation/Forth/case-insensitive.forth |   1 +
 test/Translation/Forth/comparison-ops.forth   |   1 +
 test/Translation/Forth/control-flow.forth     |   1 +
 test/Translation/Forth/do-loop.forth          |   1 +
 .../Forth/exit-outside-word-error.forth       |   1 +
 test/Translation/Forth/exit.forth             |   1 +
 test/Translation/Forth/gpu-indexing.forth     |   1 +
 .../Forth/interleaved-control-flow.forth      |   1 +
 .../Translation/Forth/leave-conditional.forth |   1 +
 test/Translation/Forth/leave.forth            |   1 +
 .../Forth/loop-index-depth-error.forth        |   1 +
 test/Translation/Forth/memory-ops.forth       |   1 +
 test/Translation/Forth/name-mangling.forth    |   1 +
 .../Forth/nested-control-flow.forth           |   1 +
 .../Forth/param-declarations.forth            |   5 +-
 .../Forth/param-ref-in-word-error.forth       |   3 +-
 .../Forth/plus-loop-negative.forth            |   1 +
 .../Forth/plus-loop-without-do-error.forth    |   1 +
 test/Translation/Forth/plus-loop.forth        |   1 +
 test/Translation/Forth/scalar-param.forth     |   8 +
 test/Translation/Forth/stack-ops.forth        |   1 +
 test/Translation/Forth/unloop-exit.forth      |   1 +
 .../Forth/unloop-without-do-error.forth       |   1 +
 test/Translation/Forth/word-definitions.forth |   1 +
 tools/warpforth-runner/warpforth-runner.cpp   |  21 +-
 54 files changed, 475 insertions(+), 103 deletions(-)
 create mode 100644 test/Translation/Forth/scalar-param.forth

diff --git a/CLAUDE.md b/CLAUDE.md
index 08b345f..7feac8f 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -88,7 +88,7 @@ uv run ruff format gpu_test/
 - **Stack Type**: `!forth.stack` - untyped stack, programmer ensures type safety
 - **Operations**: All take stack as input and produce stack as output (except `forth.stack`)
 - **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !`, `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
-- **Kernel Parameters**: Declared with `PARAM <name> <size>`, each becomes a `memref<Nxi64>` function argument with `forth.param_name` attribute. Using a param name in code pushes its byte address onto the stack via `forth.param_ref`
+- **Kernel Parameters**: Declared in the `\!` header. `\! kernel <name>` is required and must appear first. `\! param <name> i64[<N>]` becomes a `memref<Nxi64>` argument; `\! param <name> i64` becomes an `i64` argument. Using a param name in code emits `forth.param_ref` (arrays push address; scalars push value).
 - **Conversion**: `!forth.stack` → `memref<256xi64>` with explicit stack pointer
 - **GPU**: Functions wrapped in `gpu.module`, `main` gets `gpu.kernel` attribute, configured with bare pointers for NVVM conversion
 - **User-defined Words**: Modeled as `func.func` with signature `(!forth.stack) -> !forth.stack`, called via `func.call`
diff --git a/gpu_test/conftest.py b/gpu_test/conftest.py
index 3549dad..848475e 100644
--- a/gpu_test/conftest.py
+++ b/gpu_test/conftest.py
@@ -273,16 +273,66 @@ def scp_upload(self, local_path: str | Path, remote_path: str) -> None:
         )
 
 
+def _parse_array_type(type_spec: str) -> int:
+    if not type_spec.endswith("]"):
+        msg = f"Invalid array type spec: {type_spec}"
+        raise ValueError(msg)
+    base, size_str = type_spec[:-1].split("[", 1)
+    if base.lower() != "i64":
+        msg = f"Unsupported base type: {base}"
+        raise ValueError(msg)
+    return int(size_str)
+
+
+def _parse_kernel_name(forth_source: str) -> str:
+    """Parse '\\! kernel <name>' from Forth source header."""
+    for line in forth_source.splitlines():
+        stripped = line.strip()
+        if not stripped.startswith("\\!"):
+            continue
+        directive = stripped[2:].strip()
+        if "--" in directive:
+            directive = directive.split("--", 1)[0].strip()
+        if not directive:
+            continue
+        parts = directive.split()
+        if parts and parts[0].lower() == "kernel":
+            if len(parts) < 2:
+                msg = "Invalid header line: expected '\\! kernel <name>'"
+                raise ValueError(msg)
+            return parts[1]
+    msg = "Forth source has no '\\! kernel' declaration"
+    raise ValueError(msg)
+
+
 def _parse_param_declarations(forth_source: str) -> list[tuple[str, int]]:
-    """Parse 'param <name> <size>' declarations from Forth source.
+    """Parse '\\! param <name> <type>' declarations from Forth source.
 
-    Returns list of (name, size) in declaration order.
+    Returns list of (name, size) for array params in declaration order.
+    Scalar params are not supported by the GPU runner.
     """
     decls = []
     for line in forth_source.splitlines():
-        parts = line.split()
-        if len(parts) >= 3 and parts[0].upper() == "PARAM":
-            decls.append((parts[1], int(parts[2])))
+        stripped = line.strip()
+        if not stripped.startswith("\\!"):
+            continue
+        directive = stripped[2:].strip()
+        if "--" in directive:
+            directive = directive.split("--", 1)[0].strip()
+        if not directive:
+            continue
+        parts = directive.split()
+        if not parts or parts[0].lower() != "param":
+            continue
+        if len(parts) < 3:
+            msg = "Invalid header line: expected '\\! param <name> <type>'"
+            raise ValueError(msg)
+        name = parts[1]
+        type_spec = parts[2]
+        if "[" not in type_spec:
+            msg = "Scalar params are not supported by the GPU runner yet"
+            raise ValueError(msg)
+        decls.append((name, _parse_array_type(type_spec)))
     return decls
 
 
@@ -308,7 +358,8 @@ def run(
         The params dict maps param names to initial values (padded with zeros to the
         declared size). Params not in the dict are zero-initialized.
         """
-        # Parse param declarations to determine buffer sizes
+        # Parse kernel name and param declarations
+        kernel_name = _parse_kernel_name(forth_source)
         decls = _parse_param_declarations(forth_source)
         if not decls:
             msg = "Forth source has no 'param' declarations"
@@ -330,7 +381,12 @@ def run(
             ptx_path.unlink()
 
         # Build remote command
-        cmd_parts = [f"{REMOTE_TMP}/warpforth-runner", f"{REMOTE_TMP}/kernel.ptx"]
+        cmd_parts = [
+            f"{REMOTE_TMP}/warpforth-runner",
+            f"{REMOTE_TMP}/kernel.ptx",
+            "--kernel",
+            kernel_name,
+        ]
 
         for name, size in decls:
             values = params.get(name, [])
diff --git a/gpu_test/test_kernels.py b/gpu_test/test_kernels.py
index dc0bdf7..ad02a5e 100644
--- a/gpu_test/test_kernels.py
+++ b/gpu_test/test_kernels.py
@@ -18,7 +18,7 @@
 def test_addition(kernel_runner: KernelRunner) -> None:
     """3 + 4 = 7."""
     result = kernel_runner.run(
-        forth_source="PARAM DATA 256\n3 4 +\n0 CELLS DATA + !",
+        forth_source="\\! kernel main\n\\! param DATA i64[256]\n3 4 +\n0 CELLS DATA + !",
     )
     assert result[0] == 7
 
@@ -26,7 +26,7 @@ def test_addition(kernel_runner: KernelRunner) -> None:
 def test_subtraction(kernel_runner: KernelRunner) -> None:
     """10 - 3 = 7."""
     result = kernel_runner.run(
-        forth_source="PARAM DATA 256\n10 3 -\n0 CELLS DATA + !",
+        forth_source="\\! kernel main\n\\! param DATA i64[256]\n10 3 -\n0 CELLS DATA + !",
     )
     assert result[0] == 7
 
@@ -34,7 +34,7 @@ def test_subtraction(kernel_runner: KernelRunner) -> None:
 def test_multiplication(kernel_runner: KernelRunner) -> None:
     """6 * 7 = 42."""
     result = kernel_runner.run(
-        forth_source="PARAM DATA 256\n6 7 *\n0 CELLS DATA + !",
+        forth_source="\\! kernel main\n\\! param DATA i64[256]\n6 7 *\n0 CELLS DATA + !",
     )
     assert result[0] == 42
 
@@ -42,7 +42,7 @@ def test_multiplication(kernel_runner: KernelRunner) -> None:
 def test_division(kernel_runner: KernelRunner) -> None:
     """42 / 6 = 7."""
     result = kernel_runner.run(
-        forth_source="PARAM DATA 256\n42 6 /\n0 CELLS DATA + !",
+        forth_source="\\! kernel main\n\\! param DATA i64[256]\n42 6 /\n0 CELLS DATA + !",
     )
     assert result[0] == 7
 
@@ -50,7 +50,7 @@ def test_division(kernel_runner: KernelRunner) -> None:
 def test_modulo(kernel_runner: KernelRunner) -> None:
     """17 MOD 5 = 2."""
     result = kernel_runner.run(
-        forth_source="PARAM DATA 256\n17 5 MOD\n0 CELLS DATA + !",
+        forth_source="\\! kernel main\n\\! param DATA i64[256]\n17 5 MOD\n0 CELLS DATA + !",
     )
     assert result[0] == 2
 
@@ -61,7 +61,9 @@ def test_modulo(kernel_runner: KernelRunner) -> None:
 def test_dup(kernel_runner: KernelRunner) -> None:
     """DUP duplicates top of stack: 5 DUP → [5, 5]."""
     result = kernel_runner.run(
-        forth_source=("PARAM DATA 256\n5 DUP\n1 CELLS DATA + !\n0 CELLS DATA + !"),
+        forth_source=(
+            "\\! kernel main\n\\! param DATA i64[256]\n5 DUP\n1 CELLS DATA + !\n0 CELLS DATA + !"
+        ),
         output_count=2,
     )
     assert result == [5, 5]
@@ -70,7 +72,9 @@ def test_dup(kernel_runner: KernelRunner) -> None:
 def test_swap(kernel_runner: KernelRunner) -> None:
     """SWAP exchanges top two: 1 2 SWAP → [2, 1]."""
     result = kernel_runner.run(
-        forth_source=("PARAM DATA 256\n1 2 SWAP\n1 CELLS DATA + !\n0 CELLS DATA + !"),
+        forth_source=(
+            "\\! kernel main\n\\! param DATA i64[256]\n1 2 SWAP\n1 CELLS DATA + !\n0 CELLS DATA + !"
+        ),
         output_count=2,
     )
     assert result == [2, 1]
@@ -80,7 +84,12 @@ def test_over(kernel_runner: KernelRunner) -> None:
     """OVER copies second element: 1 2 OVER → [1, 2, 1]."""
     result = kernel_runner.run(
         forth_source=(
-            "PARAM DATA 256\n1 2 OVER\n2 CELLS DATA + !\n1 CELLS DATA + !\n0 CELLS DATA + !"
+            "\\! kernel main\n"
+            "\\! param DATA i64[256]\n"
+            "1 2 OVER\n"
+            "2 CELLS DATA + !\n"
+            "1 CELLS DATA + !\n"
+            "0 CELLS DATA + !"
         ),
         output_count=3,
     )
@@ -91,7 +100,12 @@ def test_rot(kernel_runner: KernelRunner) -> None:
     """ROT rotates top three: 1 2 3 ROT → [2, 3, 1]."""
     result = kernel_runner.run(
         forth_source=(
-            "PARAM DATA 256\n1 2 3 ROT\n2 CELLS DATA + !\n1 CELLS DATA + !\n0 CELLS DATA + !"
+            "\\! kernel main\n"
+            "\\! param DATA i64[256]\n"
+            "1 2 3 ROT\n"
+            "2 CELLS DATA + !\n"
+            "1 CELLS DATA + !\n"
+            "0 CELLS DATA + !"
         ),
         output_count=3,
     )
@@ -101,7 +115,7 @@ def test_rot(kernel_runner: KernelRunner) -> None:
 def test_drop(kernel_runner: KernelRunner) -> None:
     """DROP removes top: 1 2 DROP → [1]."""
     result = kernel_runner.run(
-        forth_source=("PARAM DATA 256\n1 2 DROP\n0 CELLS DATA + !"),
+        forth_source=("\\! kernel main\n\\! param DATA i64[256]\n1 2 DROP\n0 CELLS DATA + !"),
     )
     assert result[0] == 1
 
@@ -113,7 +127,7 @@ def test_comparisons(kernel_runner: KernelRunner) -> None:
     """Test =, <, >, 0= in a single kernel. True = -1, False = 0."""
     result = kernel_runner.run(
         forth_source=(
-            "PARAM DATA 256\n"
+            "\\! kernel main\n\\! param DATA i64[256]\n"
             "5 5 =  0 CELLS DATA + !\n"
             "3 5 <  1 CELLS DATA + !\n"
             "5 3 >  2 CELLS DATA + !\n"
@@ -130,7 +144,14 @@ def test_comparisons(kernel_runner: KernelRunner) -> None:
 def test_if_else_then(kernel_runner: KernelRunner) -> None:
     """IF/ELSE/THEN: if DATA[0] > 0, write 1 to DATA[1], else write 2."""
     result = kernel_runner.run(
-        forth_source=("PARAM DATA 256\n0 CELLS DATA + @\n0 >\nIF 1 ELSE 2 THEN\n1 CELLS DATA + !"),
+        forth_source=(
+            "\\! kernel main\n"
+            "\\! param DATA i64[256]\n"
+            "0 CELLS DATA + @\n"
+            "0 >\n"
+            "IF 1 ELSE 2 THEN\n"
+            "1 CELLS DATA + !"
+        ),
         params={"DATA": [5]},
         output_count=2,
     )
@@ -140,7 +161,9 @@ def test_if_else_then(kernel_runner: KernelRunner) -> None:
 def test_begin_until(kernel_runner: KernelRunner) -> None:
     """BEGIN/UNTIL countdown: 10 BEGIN 1- DUP 0= UNTIL → final value is 0."""
     result = kernel_runner.run(
-        forth_source=("PARAM DATA 256\n10 BEGIN 1 - DUP 0= UNTIL\n0 CELLS DATA + !"),
+        forth_source=(
+            "\\! kernel main\n\\! param DATA i64[256]\n10 BEGIN 1 - DUP 0= UNTIL\n0 CELLS DATA + !"
+        ),
     )
     assert result[0] == 0
 
@@ -148,7 +171,9 @@ def test_begin_until(kernel_runner: KernelRunner) -> None:
 def test_do_loop(kernel_runner: KernelRunner) -> None:
     """DO/LOOP: write I values 0..4 to DATA[0..4]."""
     result = kernel_runner.run(
-        forth_source=("PARAM DATA 256\n5 0 DO\n  I I CELLS DATA + !\nLOOP"),
+        forth_source=(
+            "\\! kernel main\n\\! param DATA i64[256]\n5 0 DO\n  I I CELLS DATA + !\nLOOP"
+        ),
         output_count=5,
     )
     assert result == [0, 1, 2, 3, 4]
@@ -157,7 +182,16 @@ def test_do_loop(kernel_runner: KernelRunner) -> None:
 def test_do_plus_loop(kernel_runner: KernelRunner) -> None:
     """DO/+LOOP: write I values 0, 2, 4, 6, 8 to DATA[0..4]."""
     result = kernel_runner.run(
-        forth_source=("PARAM DATA 256\n0\n10 0 DO\n  I OVER CELLS DATA + !\n  1 +\n2 +LOOP\nDROP"),
+        forth_source=(
+            "\\! kernel main\n"
+            "\\! param DATA i64[256]\n"
+            "0\n"
+            "10 0 DO\n"
+            "  I OVER CELLS DATA + !\n"
+            "  1 +\n"
+            "2 +LOOP\n"
+            "DROP"
+        ),
         output_count=5,
     )
     assert result == [0, 2, 4, 6, 8]
@@ -166,7 +200,16 @@ def test_do_plus_loop(kernel_runner: KernelRunner) -> None:
 def test_do_plus_loop_negative(kernel_runner: KernelRunner) -> None:
     """DO/+LOOP with negative step: count down from 10 to 1."""
     result = kernel_runner.run(
-        forth_source=("PARAM DATA 256\n0\n0 10 DO\n  I OVER CELLS DATA + !\n  1 +\n-1 +LOOP\nDROP"),
+        forth_source=(
+            "\\! kernel main\n"
+            "\\! param DATA i64[256]\n"
+            "0\n"
+            "0 10 DO\n"
+            "  I OVER CELLS DATA + !\n"
+            "  1 +\n"
+            "-1 +LOOP\n"
+            "DROP"
+        ),
         output_count=10,
     )
     assert result == [10, 9, 8, 7, 6, 5, 4, 3, 2, 1]
@@ -180,7 +223,7 @@ def test_multi_while(kernel_runner: KernelRunner) -> None:
     """
     result = kernel_runner.run(
         forth_source=(
-            "PARAM DATA 256\n"
+            "\\! kernel main\n\\! param DATA i64[256]\n"
             "20 BEGIN DUP 10 > WHILE DUP 2 MOD 0= WHILE 1 - REPEAT THEN\n"
             "0 CELLS DATA + !"
         ),
@@ -196,7 +239,10 @@ def test_while_until(kernel_runner: KernelRunner) -> None:
     """
     result = kernel_runner.run(
         forth_source=(
-            "PARAM DATA 256\n10 BEGIN DUP 0 > WHILE 1 - DUP 5 = UNTIL THEN\n0 CELLS DATA + !"
+            "\\! kernel main\n"
+            "\\! param DATA i64[256]\n"
+            "10 BEGIN DUP 0 > WHILE 1 - DUP 5 = UNTIL THEN\n"
+            "0 CELLS DATA + !"
         ),
     )
     assert result[0] == 5
@@ -208,7 +254,7 @@ def test_while_until(kernel_runner: KernelRunner) -> None:
 def test_global_id(kernel_runner: KernelRunner) -> None:
     """4 threads each write GLOBAL-ID to DATA[GLOBAL-ID]."""
     result = kernel_runner.run(
-        forth_source=("PARAM DATA 256\nGLOBAL-ID\nDUP CELLS DATA + !"),
+        forth_source=("\\! kernel main\n\\! param DATA i64[256]\nGLOBAL-ID\nDUP CELLS DATA + !"),
         block=(4, 1, 1),
         output_count=4,
     )
@@ -219,8 +265,8 @@ def test_multi_param(kernel_runner: KernelRunner) -> None:
     """Two params: each thread reads INPUT[i], doubles it, writes OUTPUT[i]."""
     result = kernel_runner.run(
         forth_source=(
-            "PARAM INPUT 4\n"
-            "PARAM OUTPUT 4\n"
+            "\\! kernel main\n\\! param INPUT i64[4]\n"
+            "\\! param OUTPUT i64[4]\n"
             "GLOBAL-ID\n"
             "DUP CELLS INPUT + @\n"
             "DUP +\n"
@@ -243,9 +289,9 @@ def test_naive_matmul_i64(kernel_runner: KernelRunner) -> None:
     # GLOBAL-ID maps to (row, col) with row = gid / N, col = gid MOD N.
     result = kernel_runner.run(
         forth_source=(
-            "PARAM A 8\n"
-            "PARAM B 12\n"
-            "PARAM C 6\n"
+            "\\! kernel main\n\\! param A i64[8]\n"
+            "\\! param B i64[12]\n"
+            "\\! param C i64[6]\n"
             "GLOBAL-ID\n"
             "DUP 3 /\n"
             "SWAP 3 MOD\n"
@@ -277,6 +323,8 @@ def test_naive_matmul_i64(kernel_runner: KernelRunner) -> None:
 def test_user_defined_word(kernel_runner: KernelRunner) -> None:
     """: DOUBLE DUP + ; then 5 DOUBLE → 10."""
     result = kernel_runner.run(
-        forth_source=("PARAM DATA 256\n: DOUBLE DUP + ;\n5 DOUBLE\n0 CELLS DATA + !"),
+        forth_source=(
+            "\\! kernel main\n\\! param DATA i64[256]\n: DOUBLE DUP + ;\n5 DOUBLE\n0 CELLS DATA + !"
+        ),
     )
     assert result[0] == 10
diff --git a/lib/Conversion/ForthToGPU/ForthToGPU.cpp b/lib/Conversion/ForthToGPU/ForthToGPU.cpp
index 043f66d..7d72597 100644
--- a/lib/Conversion/ForthToGPU/ForthToGPU.cpp
+++ b/lib/Conversion/ForthToGPU/ForthToGPU.cpp
@@ -182,7 +182,7 @@ struct ConvertForthToGPUPass
 
   void convertFuncToGPU(func::FuncOp funcOp, gpu::GPUModuleOp gpuModule,
                         IRRewriter &rewriter) {
-    bool isKernel = funcOp.getName() == "main";
+    bool isKernel = funcOp->hasAttr("forth.kernel");
 
     if (isKernel) {
       auto gpuFunc = createGPUFunc(funcOp, gpuModule, rewriter);
diff --git a/lib/Conversion/ForthToMemRef/ForthToMemRef.cpp b/lib/Conversion/ForthToMemRef/ForthToMemRef.cpp
index 9a70304..f0031d1 100644
--- a/lib/Conversion/ForthToMemRef/ForthToMemRef.cpp
+++ b/lib/Conversion/ForthToMemRef/ForthToMemRef.cpp
@@ -632,14 +632,23 @@ struct ParamRefOpConversion : public OpConversionPattern<forth::ParamRefOp> {
       return rewriter.notifyMatchFailure(
           op, "no function argument with param_name: " + paramName);
 
-    // Extract pointer as index, then cast to i64
-    Value ptrIndex =
-        rewriter.create<memref::ExtractAlignedPointerAsIndexOp>(loc, memrefArg);
-    Value ptrI64 = rewriter.create<arith::IndexCastOp>(
-        loc, rewriter.getI64Type(), ptrIndex);
+    Value valueToPush;
+    if (auto memrefType = dyn_cast<MemRefType>(memrefArg.getType())) {
+      // Extract pointer as index, then cast to i64
+      Value ptrIndex = rewriter.create<memref::ExtractAlignedPointerAsIndexOp>(
+          loc, memrefArg);
+      valueToPush = rewriter.create<arith::IndexCastOp>(
+          loc, rewriter.getI64Type(), ptrIndex);
+    } else if (memrefArg.getType().isInteger(64)) {
+      // Scalar param: push value directly.
+      valueToPush = memrefArg;
+    } else {
+      return rewriter.notifyMatchFailure(
+          op, "unsupported param argument type for param_ref");
+    }
 
     // Push onto stack
-    Value newSP = pushValue(loc, rewriter, memref, stackPtr, ptrI64);
+    Value newSP = pushValue(loc, rewriter, memref, stackPtr, valueToPush);
 
     rewriter.replaceOpWithMultiple(op, {{memref, newSP}});
     return success();
diff --git a/lib/Translation/ForthToMLIR/ForthToMLIR.cpp b/lib/Translation/ForthToMLIR/ForthToMLIR.cpp
index fb992b2..952faa0 100644
--- a/lib/Translation/ForthToMLIR/ForthToMLIR.cpp
+++ b/lib/Translation/ForthToMLIR/ForthToMLIR.cpp
@@ -16,6 +16,7 @@
 #include "warpforth/Dialect/Forth/ForthDialect.h"
 #include "warpforth/Translation/ForthToMLIR/ForthToMLIR.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include <cctype>
 
@@ -107,6 +108,12 @@ void ForthLexer::reset() {
   endPtr = buffer->getBufferEnd();
 }
 
+void ForthLexer::resetTo(const char *ptr) {
+  auto buffer = sourceMgr.getMemoryBuffer(bufferID);
+  curPtr = ptr;
+  endPtr = buffer->getBufferEnd();
+}
+
 bool ForthLexer::isNumber(const std::string &str) const {
   if (str.empty())
     return false;
@@ -185,26 +192,202 @@ LogicalResult ForthParser::emitError(const llvm::Twine &message) {
   return failure();
 }
 
-void ForthParser::scanParamDeclarations() {
-  lexer.reset();
-  consume();
-  while (currentToken.kind != Token::Kind::EndOfFile) {
-    if (currentToken.kind == Token::Kind::Word &&
-        currentToken.text == "PARAM") {
-      consume(); // consume "PARAM"
-      if (currentToken.kind != Token::Kind::Word)
-        continue;
-      std::string name = currentToken.text;
-      consume(); // consume name
-      if (currentToken.kind != Token::Kind::Number)
-        continue;
-      int64_t size = std::stoll(currentToken.text);
-      consume(); // consume size
-      paramDecls.push_back({name, size});
+LogicalResult ForthParser::emitErrorAt(llvm::SMLoc loc,
+                                       const llvm::Twine &message) {
+  sourceMgr.PrintMessage(loc, llvm::SourceMgr::DK_Error, message);
+  return failure();
+}
+
+LogicalResult ForthParser::parseHeader() {
+  auto buffer = sourceMgr.getMemoryBuffer(sourceMgr.getMainFileID());
+  const char *bufStart = buffer->getBufferStart();
+  const char *bufEnd = buffer->getBufferEnd();
+
+  auto ltrim = [](llvm::StringRef s) {
+    while (!s.empty() && (s.front() == ' ' || s.front() == '\t'))
+      s = s.drop_front();
+    return s;
+  };
+  auto rtrim = [](llvm::StringRef s) {
+    while (!s.empty() &&
+           (s.back() == ' ' || s.back() == '\t' || s.back() == '\r'))
+      s = s.drop_back();
+    return s;
+  };
+  auto trim = [&](llvm::StringRef s) { return rtrim(ltrim(s)); };
+
+  auto splitWS = [](llvm::StringRef s) {
+    SmallVector<llvm::StringRef> parts;
+    while (!s.empty()) {
+      while (!s.empty() && (s.front() == ' ' || s.front() == '\t'))
+        s = s.drop_front();
+      if (s.empty())
+        break;
+      size_t i = 0;
+      while (i < s.size() && s[i] != ' ' && s[i] != '\t')
+        ++i;
+      parts.push_back(s.substr(0, i));
+      s = s.substr(i);
+    }
+    return parts;
+  };
+
+  bool headerEnded = false;
+  bool sawKernel = false;
+  bool sawDirective = false;
+  headerEndPtr = bufStart;
+
+  const char *lineStart = bufStart;
+  while (lineStart <= bufEnd) {
+    const char *lineEnd = lineStart;
+    while (lineEnd < bufEnd && *lineEnd != '\n' && *lineEnd != '\r')
+      ++lineEnd;
+    llvm::StringRef line(lineStart, lineEnd - lineStart);
+    llvm::StringRef trimmed = trim(line);
+
+    auto lineLoc = llvm::SMLoc::getFromPointer(lineStart);
+
+    if (trimmed.empty()) {
+      if (!headerEnded)
+        headerEndPtr = lineEnd < bufEnd ? lineEnd + 1 : lineEnd;
+    } else if (trimmed.starts_with("\\!")) {
+      if (headerEnded) {
+        return emitErrorAt(lineLoc,
+                           "header directive must appear before any code");
+      }
+
+      llvm::StringRef directiveLine = trim(trimmed.drop_front(2));
+      size_t annotPos = directiveLine.find("--");
+      if (annotPos != llvm::StringRef::npos) {
+        directiveLine = trim(directiveLine.substr(0, annotPos));
+      }
+
+      if (directiveLine.empty()) {
+        return emitErrorAt(lineLoc, "empty header directive");
+      }
+
+      auto tokens = splitWS(directiveLine);
+      if (tokens.empty()) {
+        return emitErrorAt(lineLoc, "empty header directive");
+      }
+
+      std::string directive = toUpperCase(tokens[0]);
+      if (!sawDirective && directive != "KERNEL") {
+        return emitErrorAt(lineLoc, "\\! kernel must appear first");
+      }
+
+      if (directive == "KERNEL") {
+        if (sawKernel) {
+          return emitErrorAt(lineLoc, "duplicate \\! kernel directive");
+        }
+        if (tokens.size() != 2) {
+          return emitErrorAt(lineLoc,
+                             "kernel directive expects: \\! kernel <name>");
+        }
+        kernelName = tokens[1].str();
+        sawKernel = true;
+        sawDirective = true;
+      } else if (directive == "PARAM" || directive == "SHARED") {
+        if (!sawKernel) {
+          return emitErrorAt(lineLoc, "\\! kernel must appear first");
+        }
+        if (tokens.size() != 3) {
+          return emitErrorAt(lineLoc,
+                             "param/shared directive expects: \\! param "
+                             "<name> <type>");
+        }
+
+        std::string nameUpper = toUpperCase(tokens[1]);
+        for (const auto &param : paramDecls) {
+          if (param.name == nameUpper) {
+            return emitErrorAt(lineLoc,
+                               "duplicate parameter name: " + nameUpper);
+          }
+        }
+        for (const auto &shared : sharedDecls) {
+          if (shared.name == nameUpper) {
+            return emitErrorAt(lineLoc, "duplicate shared name: " + nameUpper);
+          }
+        }
+
+        llvm::StringRef typeToken = tokens[2];
+        bool isArray = false;
+        int64_t size = 0;
+        size_t lbracket = typeToken.find('[');
+        if (lbracket != llvm::StringRef::npos) {
+          size_t rbracket = typeToken.find(']');
+          if (rbracket == llvm::StringRef::npos ||
+              rbracket != typeToken.size() - 1) {
+            return emitErrorAt(lineLoc,
+                               "array type must use suffix [N], e.g. i64[4]");
+          }
+          llvm::StringRef base = typeToken.substr(0, lbracket);
+          llvm::StringRef sizeStr =
+              typeToken.substr(lbracket + 1, rbracket - lbracket - 1);
+          if (toUpperCase(base) != "I64") {
+            return emitErrorAt(lineLoc, "unsupported base type: " + base.str());
+          }
+          if (sizeStr.empty())
+            return emitErrorAt(lineLoc, "array type requires a size");
+          for (char c : sizeStr) {
+            if (!std::isdigit(static_cast<unsigned char>(c)))
+              return emitErrorAt(lineLoc, "array size must be an integer");
+          }
+          size = std::stoll(sizeStr.str());
+          if (size <= 0)
+            return emitErrorAt(lineLoc, "array size must be positive");
+          isArray = true;
+        } else {
+          if (toUpperCase(typeToken) != "I64") {
+            return emitErrorAt(lineLoc,
+                               "unsupported scalar type: " + typeToken.str());
+          }
+        }
+
+        if (directive == "PARAM") {
+          ParamDecl decl;
+          decl.name = nameUpper;
+          decl.isArray = isArray;
+          decl.size = size;
+          paramDecls.push_back(decl);
+        } else {
+          SharedDecl decl;
+          decl.name = nameUpper;
+          decl.isArray = isArray;
+          decl.size = size;
+          sharedDecls.push_back(decl);
+        }
+        sawDirective = true;
+      } else {
+        return emitErrorAt(lineLoc, "unknown header directive: " + directive);
+      }
+
+      headerEndPtr = lineEnd < bufEnd ? lineEnd + 1 : lineEnd;
+    } else if (trimmed.starts_with("\\")) {
+      if (trimmed.size() == 1 || trimmed[1] == ' ' || trimmed[1] == '\t') {
+        if (!headerEnded)
+          headerEndPtr = lineEnd < bufEnd ? lineEnd + 1 : lineEnd;
+      } else {
+        headerEnded = true;
+      }
     } else {
-      consume();
+      headerEnded = true;
     }
+
+    if (lineEnd >= bufEnd)
+      break;
+    if (*lineEnd == '\r' && lineEnd + 1 < bufEnd && lineEnd[1] == '\n')
+      lineStart = lineEnd + 2;
+    else
+      lineStart = lineEnd + 1;
+  }
+
+  if (!sawKernel) {
+    auto loc = llvm::SMLoc::getFromPointer(bufStart);
+    return emitErrorAt(loc, "\\! kernel <name> is required");
   }
+
+  return success();
 }
 
 Value ForthParser::emitOperation(StringRef word, Value inputStack,
@@ -434,15 +617,6 @@ LogicalResult ForthParser::parseBody(Value &stack) {
          currentToken.kind != Token::Kind::Semicolon &&
          currentToken.kind != Token::Kind::Colon) {
 
-    // Skip PARAM declarations at top level.
-    if (!inWordDefinition && currentToken.kind == Token::Kind::Word &&
-        currentToken.text == "PARAM") {
-      consume(); // "PARAM"
-      consume(); // name
-      consume(); // size
-      continue;
-    }
-
     if (currentToken.kind == Token::Kind::Number) {
       Location tokenLoc = getLoc();
       int64_t value = std::stoll(currentToken.text);
@@ -823,11 +997,11 @@ OwningOpRef<ModuleOp> ForthParser::parseModule() {
 
   builder.setInsertionPointToEnd(module->getBody());
 
-  // Pre-pass: scan for param declarations
-  scanParamDeclarations();
+  if (failed(parseHeader()))
+    return nullptr;
 
   // First pass: parse all word definitions
-  lexer.reset();
+  lexer.resetTo(headerEndPtr);
   consume();
   while (currentToken.kind != Token::Kind::EndOfFile) {
     if (currentToken.kind == Token::Kind::Colon) {
@@ -840,7 +1014,7 @@ OwningOpRef<ModuleOp> ForthParser::parseModule() {
   }
 
   // Reset lexer for second pass
-  lexer.reset();
+  lexer.resetTo(headerEndPtr);
   consume();
 
   // Reset insertion point to end of module for main function
@@ -849,12 +1023,17 @@ OwningOpRef<ModuleOp> ForthParser::parseModule() {
   // Build function argument types from param declarations
   SmallVector<Type> argTypes;
   for (const auto &param : paramDecls) {
-    argTypes.push_back(MemRefType::get({param.size}, builder.getI64Type()));
+    if (param.isArray) {
+      argTypes.push_back(MemRefType::get({param.size}, builder.getI64Type()));
+    } else {
+      argTypes.push_back(builder.getI64Type());
+    }
   }
 
   auto funcType = builder.getFunctionType(argTypes, {});
-  auto funcOp = builder.create<func::FuncOp>(loc, "main", funcType);
+  auto funcOp = builder.create<func::FuncOp>(loc, kernelName, funcType);
   funcOp.setPrivate();
+  funcOp->setAttr("forth.kernel", builder.getUnitAttr());
 
   // Annotate arguments with param names
   for (size_t i = 0; i < paramDecls.size(); ++i) {
diff --git a/lib/Translation/ForthToMLIR/ForthToMLIR.h b/lib/Translation/ForthToMLIR/ForthToMLIR.h
index c860dd2..46850da 100644
--- a/lib/Translation/ForthToMLIR/ForthToMLIR.h
+++ b/lib/Translation/ForthToMLIR/ForthToMLIR.h
@@ -17,10 +17,18 @@
 namespace mlir {
 namespace forth {
 
-/// A declared kernel parameter: `param <name> <size>`.
+/// A declared kernel parameter: `param <name> <type>`.
 struct ParamDecl {
   std::string name;
-  int64_t size;
+  bool isArray = false;
+  int64_t size = 0;
+};
+
+/// A declared shared memory region: `shared <name> <type>`.
+struct SharedDecl {
+  std::string name;
+  bool isArray = false;
+  int64_t size = 0;
 };
 
 /// Simple token representing a Forth word or literal.
@@ -47,6 +55,9 @@ class ForthLexer {
   /// Reset lexer to beginning of buffer.
   void reset();
 
+  /// Reset lexer to a specific position in the buffer.
+  void resetTo(const char *ptr);
+
 private:
   llvm::SourceMgr &sourceMgr;
   unsigned bufferID;
@@ -79,6 +90,9 @@ class ForthParser {
   Token currentToken;
   std::unordered_set<std::string> wordDefs;
   std::vector<ParamDecl> paramDecls;
+  std::vector<SharedDecl> sharedDecls;
+  std::string kernelName;
+  const char *headerEndPtr = nullptr;
   bool inWordDefinition = false;
 
   /// Control flow stack tag: Orig = forward reference (IF->THEN, WHILE->exit),
@@ -97,8 +111,8 @@ class ForthParser {
   };
   SmallVector<LoopContext> loopStack;
 
-  /// Scan for `param <name> <size>` declarations (pre-pass).
-  void scanParamDeclarations();
+  /// Parse the kernel header (\\! directives) and record metadata.
+  LogicalResult parseHeader();
 
   /// Advance to the next token.
   void consume();
@@ -106,6 +120,9 @@ class ForthParser {
   /// Emit an error at the current location.
   LogicalResult emitError(const llvm::Twine &message);
 
+  /// Emit an error at a specific buffer location.
+  LogicalResult emitErrorAt(llvm::SMLoc loc, const llvm::Twine &message);
+
   /// Parse a sequence of Forth operations.
   LogicalResult parseOperations(Value &stack);
 
diff --git a/test/Conversion/ForthToGPU/basic-gpu-wrap.mlir b/test/Conversion/ForthToGPU/basic-gpu-wrap.mlir
index f78e7d5..5e30344 100644
--- a/test/Conversion/ForthToGPU/basic-gpu-wrap.mlir
+++ b/test/Conversion/ForthToGPU/basic-gpu-wrap.mlir
@@ -6,7 +6,7 @@
 // CHECK: gpu.return
 
 module {
-  func.func private @main() {
+  func.func private @main() attributes {forth.kernel} {
     %alloca = memref.alloca() : memref<256xi64>
     %c0 = arith.constant 0 : index
     return
diff --git a/test/Conversion/ForthToGPU/intrinsic-conversion.mlir b/test/Conversion/ForthToGPU/intrinsic-conversion.mlir
index 42d245b..0c94ed7 100644
--- a/test/Conversion/ForthToGPU/intrinsic-conversion.mlir
+++ b/test/Conversion/ForthToGPU/intrinsic-conversion.mlir
@@ -10,7 +10,7 @@
 // CHECK: gpu.grid_dim y
 
 module {
-  func.func private @main() {
+  func.func private @main() attributes {forth.kernel} {
     %alloca = memref.alloca() : memref<256xi64>
     %c0 = arith.constant 0 : index
     %c1 = arith.constant 1 : index
diff --git a/test/Conversion/ForthToGPU/private-functions.mlir b/test/Conversion/ForthToGPU/private-functions.mlir
index 5e19560..ed22b83 100644
--- a/test/Conversion/ForthToGPU/private-functions.mlir
+++ b/test/Conversion/ForthToGPU/private-functions.mlir
@@ -14,7 +14,7 @@ module {
   func.func private @helper(%arg0: memref<256xi64>, %arg1: index) -> (memref<256xi64>, index) {
     return %arg0, %arg1 : memref<256xi64>, index
   }
-  func.func private @main() {
+  func.func private @main() attributes {forth.kernel} {
     %alloca = memref.alloca() : memref<256xi64>
     %c0 = arith.constant 0 : index
     return
diff --git a/test/Pipeline/begin-until.forth b/test/Pipeline/begin-until.forth
index bd5aa70..189b487 100644
--- a/test/Pipeline/begin-until.forth
+++ b/test/Pipeline/begin-until.forth
@@ -11,5 +11,6 @@
 \ MID: cf.cond_br
 \ MID: gpu.return
 
-PARAM DATA 4
+\! kernel main
+\! param DATA i64[4]
 10 BEGIN 1 - DUP 0= UNTIL DATA 0 CELLS + !
diff --git a/test/Pipeline/begin-while-repeat.forth b/test/Pipeline/begin-while-repeat.forth
index b19b9c8..b2998fe 100644
--- a/test/Pipeline/begin-while-repeat.forth
+++ b/test/Pipeline/begin-while-repeat.forth
@@ -11,5 +11,6 @@
 \ MID: cf.cond_br
 \ MID: gpu.return
 
-PARAM DATA 4
+\! kernel main
+\! param DATA i64[4]
 10 BEGIN DUP 0 > WHILE 1 - REPEAT DATA 0 CELLS + !
diff --git a/test/Pipeline/control-flow.forth b/test/Pipeline/control-flow.forth
index 9b80708..6fac4d8 100644
--- a/test/Pipeline/control-flow.forth
+++ b/test/Pipeline/control-flow.forth
@@ -12,5 +12,6 @@
 \ MID: cf.cond_br
 \ MID: gpu.return
 
-PARAM DATA 256
+\! kernel main
+\! param DATA i64[256]
 DATA @ 5 > IF DATA @ 1 + DATA ! THEN
diff --git a/test/Pipeline/do-loop.forth b/test/Pipeline/do-loop.forth
index 52c5322..57fb4b8 100644
--- a/test/Pipeline/do-loop.forth
+++ b/test/Pipeline/do-loop.forth
@@ -11,5 +11,6 @@
 \ MID: cf.cond_br
 \ MID: gpu.return
 
-PARAM DATA 4
+\! kernel main
+\! param DATA i64[4]
 10 0 DO I LOOP DATA 0 CELLS + !
diff --git a/test/Pipeline/exit.forth b/test/Pipeline/exit.forth
index 4a29b93..6dfa82b 100644
--- a/test/Pipeline/exit.forth
+++ b/test/Pipeline/exit.forth
@@ -1,6 +1,7 @@
 \ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt --warpforth-pipeline | %FileCheck %s
 \ CHECK: gpu.binary @warpforth_module
 
-PARAM DATA 4
+\! kernel main
+\! param DATA i64[4]
 : DO-EXIT 1 IF EXIT THEN 42 ;
 DO-EXIT DATA 0 CELLS + !
diff --git a/test/Pipeline/full-pipeline.forth b/test/Pipeline/full-pipeline.forth
index a2cd8c5..c2c77fb 100644
--- a/test/Pipeline/full-pipeline.forth
+++ b/test/Pipeline/full-pipeline.forth
@@ -14,7 +14,8 @@
 \ MID: llvm.store
 \ MID: gpu.return
 
-PARAM DATA 256
+\! kernel main
+\! param DATA i64[256]
 GLOBAL-ID CELLS DATA + @
 1 +
 GLOBAL-ID CELLS DATA + !
diff --git a/test/Pipeline/interleaved-control-flow.forth b/test/Pipeline/interleaved-control-flow.forth
index 544b917..2b5573d 100644
--- a/test/Pipeline/interleaved-control-flow.forth
+++ b/test/Pipeline/interleaved-control-flow.forth
@@ -22,7 +22,8 @@
 \ MID: cf.cond_br
 \ MID: cf.br
 
-PARAM DATA 4
+\! kernel main
+\! param DATA i64[4]
 : multi-while
   BEGIN DUP 10 > WHILE DUP 2 MOD 0= WHILE 1 - REPEAT DROP THEN ;
 : while-until
diff --git a/test/Pipeline/leave.forth b/test/Pipeline/leave.forth
index f6bcb6f..3fbe231 100644
--- a/test/Pipeline/leave.forth
+++ b/test/Pipeline/leave.forth
@@ -11,5 +11,6 @@
 \ MID: cf.cond_br
 \ MID: gpu.return
 
-PARAM DATA 4
+\! kernel main
+\! param DATA i64[4]
 10 0 DO LEAVE LOOP DATA 0 CELLS + !
diff --git a/test/Pipeline/matmul-naive.forth b/test/Pipeline/matmul-naive.forth
index fe87d58..4cae506 100644
--- a/test/Pipeline/matmul-naive.forth
+++ b/test/Pipeline/matmul-naive.forth
@@ -7,9 +7,10 @@
 \ Verify the kernel signature at the memref+gpu stage.
 \ MID: gpu.func @main(%arg0: memref<8xi64> {forth.param_name = "A"}, %arg1: memref<12xi64> {forth.param_name = "B"}, %arg2: memref<6xi64> {forth.param_name = "C"}) kernel
 
-PARAM A 8
-PARAM B 12
-PARAM C 6
+\! kernel main
+\! param A i64[8]
+\! param B i64[12]
+\! param C i64[6]
 
 \ M=2, N=3, K=4. One thread computes C[row, col] where gid = row*N + col.
 GLOBAL-ID
diff --git a/test/Pipeline/multi-param.forth b/test/Pipeline/multi-param.forth
index 1fdfbde..3cd9b43 100644
--- a/test/Pipeline/multi-param.forth
+++ b/test/Pipeline/multi-param.forth
@@ -12,8 +12,9 @@
 \ MID: llvm.store
 \ MID: gpu.return
 
-PARAM INPUT 256
-PARAM OUTPUT 256
+\! kernel main
+\! param INPUT i64[256]
+\! param OUTPUT i64[256]
 GLOBAL-ID CELLS INPUT + @
 2 *
 GLOBAL-ID CELLS OUTPUT + !
diff --git a/test/Pipeline/nested-control-flow.forth b/test/Pipeline/nested-control-flow.forth
index 41a520b..38066f0 100644
--- a/test/Pipeline/nested-control-flow.forth
+++ b/test/Pipeline/nested-control-flow.forth
@@ -10,5 +10,6 @@
 \ MID: cf.br
 \ MID: arith.xori
 
-PARAM DATA 4
+\! kernel main
+\! param DATA i64[4]
 3 0 DO 4 0 DO J I + LOOP LOOP DATA 0 CELLS + !
diff --git a/test/Pipeline/plus-loop-negative.forth b/test/Pipeline/plus-loop-negative.forth
index 50afe63..4070e47 100644
--- a/test/Pipeline/plus-loop-negative.forth
+++ b/test/Pipeline/plus-loop-negative.forth
@@ -3,5 +3,6 @@
 \ Verify that +LOOP with negative step through the full pipeline produces a gpu.binary
 \ CHECK: gpu.binary @warpforth_module
 
-PARAM DATA 4
+\! kernel main
+\! param DATA i64[4]
 0 10 DO I DATA 0 CELLS + ! -1 +LOOP
diff --git a/test/Pipeline/plus-loop.forth b/test/Pipeline/plus-loop.forth
index 6eb32ed..87e0f9c 100644
--- a/test/Pipeline/plus-loop.forth
+++ b/test/Pipeline/plus-loop.forth
@@ -3,5 +3,6 @@
 \ Verify that +LOOP through the full pipeline produces a gpu.binary
 \ CHECK: gpu.binary @warpforth_module
 
-PARAM DATA 4
+\! kernel main
+\! param DATA i64[4]
 10 0 DO I DATA 0 CELLS + ! 2 +LOOP
diff --git a/test/Pipeline/unloop-exit.forth b/test/Pipeline/unloop-exit.forth
index 65bbc50..5de0b28 100644
--- a/test/Pipeline/unloop-exit.forth
+++ b/test/Pipeline/unloop-exit.forth
@@ -1,6 +1,7 @@
 \ RUN: %warpforth-translate --forth-to-mlir %s | %warpforth-opt --warpforth-pipeline | %FileCheck %s
 \ CHECK: gpu.binary @warpforth_module
 
-PARAM DATA 4
+\! kernel main
+\! param DATA i64[4]
 : FIND-FIVE  10 0 DO I 5 = IF UNLOOP EXIT THEN LOOP 0 ;
 FIND-FIVE DATA 0 CELLS + !
diff --git a/test/Translation/Forth/arithmetic-ops.forth b/test/Translation/Forth/arithmetic-ops.forth
index 612c526..4f767e2 100644
--- a/test/Translation/Forth/arithmetic-ops.forth
+++ b/test/Translation/Forth/arithmetic-ops.forth
@@ -9,4 +9,5 @@
 \ CHECK: %[[S5:.*]] = forth.mul %[[S4]]
 \ CHECK: %[[S6:.*]] = forth.div %[[S5]]
 \ CHECK: %{{.*}} = forth.mod %[[S6]]
+\! kernel main
 1 2 + - * / MOD
diff --git a/test/Translation/Forth/basic-literals.forth b/test/Translation/Forth/basic-literals.forth
index c77428f..9cff6b5 100644
--- a/test/Translation/Forth/basic-literals.forth
+++ b/test/Translation/Forth/basic-literals.forth
@@ -4,4 +4,5 @@
 \ CHECK-NEXT: forth.literal %{{.*}} 42
 \ CHECK-NEXT: forth.literal %{{.*}} -7
 \ CHECK-NEXT: forth.literal %{{.*}} 0
+\! kernel main
 42 -7 0
diff --git a/test/Translation/Forth/begin-until.forth b/test/Translation/Forth/begin-until.forth
index dbd2c27..fec8b73 100644
--- a/test/Translation/Forth/begin-until.forth
+++ b/test/Translation/Forth/begin-until.forth
@@ -14,4 +14,5 @@
 \ CHECK-NEXT:  cf.cond_br %[[FLAG]], ^bb2(%[[PF]] : !forth.stack), ^bb1(%[[PF]] : !forth.stack)
 \ CHECK:     ^bb2(%[[B2:.*]]: !forth.stack):
 \ CHECK-NEXT:  return
+\! kernel main
 10 BEGIN 1 - DUP 0= UNTIL
diff --git a/test/Translation/Forth/begin-while-repeat.forth b/test/Translation/Forth/begin-while-repeat.forth
index 873f44c..90a2a96 100644
--- a/test/Translation/Forth/begin-while-repeat.forth
+++ b/test/Translation/Forth/begin-while-repeat.forth
@@ -17,4 +17,5 @@
 \ CHECK-NEXT:  cf.br ^bb1(%[[SUB]] : !forth.stack)
 \ CHECK:     ^bb3(%[[B3:.*]]: !forth.stack):
 \ CHECK-NEXT:  return
+\! kernel main
 10 BEGIN DUP 0 > WHILE 1 - REPEAT
diff --git a/test/Translation/Forth/bitwise-ops.forth b/test/Translation/Forth/bitwise-ops.forth
index c2d4742..cdf7c9e 100644
--- a/test/Translation/Forth/bitwise-ops.forth
+++ b/test/Translation/Forth/bitwise-ops.forth
@@ -19,4 +19,5 @@
 \ CHECK: %[[S15:.*]] = forth.literal %[[S14]]
 \ CHECK: %[[S16:.*]] = forth.literal %[[S15]]
 \ CHECK: %{{.*}} = forth.rshift %[[S16]]
+\! kernel main
 3 5 AND 7 8 OR 15 3 XOR 42 NOT 1 4 LSHIFT 256 2 RSHIFT
diff --git a/test/Translation/Forth/case-insensitive.forth b/test/Translation/Forth/case-insensitive.forth
index b80310a..c545f15 100644
--- a/test/Translation/Forth/case-insensitive.forth
+++ b/test/Translation/Forth/case-insensitive.forth
@@ -5,4 +5,5 @@
 \ CHECK: forth.drop
 \ CHECK: forth.swap
 \ CHECK: forth.add
+\! kernel main
 1 Dup DROP swap duP +
diff --git a/test/Translation/Forth/comparison-ops.forth b/test/Translation/Forth/comparison-ops.forth
index 5e5746b..804b5bb 100644
--- a/test/Translation/Forth/comparison-ops.forth
+++ b/test/Translation/Forth/comparison-ops.forth
@@ -22,4 +22,5 @@
 \ CHECK: %[[S18:.*]] = forth.literal %[[S17]]
 \ CHECK: %[[S19:.*]] = forth.literal %[[S18]]
 \ CHECK: %{{.*}} = forth.ge %[[S19]]
+\! kernel main
 1 2 = 3 4 < 5 6 > 0 0= 7 8 <> 9 10 <= 11 12 >=
diff --git a/test/Translation/Forth/control-flow.forth b/test/Translation/Forth/control-flow.forth
index 8e03361..f8a39a2 100644
--- a/test/Translation/Forth/control-flow.forth
+++ b/test/Translation/Forth/control-flow.forth
@@ -13,6 +13,7 @@
 \ CHECK:     ^bb2(%[[B2:.*]]: !forth.stack):
 \ CHECK-NEXT:  %[[L99:.*]] = forth.literal %[[B2]] 99 : !forth.stack -> !forth.stack
 \ CHECK-NEXT:  cf.br ^bb3(%[[L99]] : !forth.stack)
+\! kernel main
 1 IF 42 ELSE 99 THEN
 
 \ Basic IF/THEN (no ELSE - fallthrough on false)
diff --git a/test/Translation/Forth/do-loop.forth b/test/Translation/Forth/do-loop.forth
index 4c0b9bc..39e8e0f 100644
--- a/test/Translation/Forth/do-loop.forth
+++ b/test/Translation/Forth/do-loop.forth
@@ -28,4 +28,5 @@
 \ CHECK-NEXT:  cf.cond_br %[[CROSSED]], ^bb2(%[[PUSH]] : !forth.stack), ^bb1(%[[PUSH]] : !forth.stack)
 \ CHECK:     ^bb2(%[[B2:.*]]: !forth.stack):
 \ CHECK-NEXT:  return
+\! kernel main
 10 0 DO I LOOP
diff --git a/test/Translation/Forth/exit-outside-word-error.forth b/test/Translation/Forth/exit-outside-word-error.forth
index 07d7fe9..410f099 100644
--- a/test/Translation/Forth/exit-outside-word-error.forth
+++ b/test/Translation/Forth/exit-outside-word-error.forth
@@ -1,3 +1,4 @@
 \ RUN: %not %warpforth-translate --forth-to-mlir %s 2>&1 | %FileCheck %s
 \ CHECK: EXIT outside word definition
+\! kernel main
 EXIT
diff --git a/test/Translation/Forth/exit.forth b/test/Translation/Forth/exit.forth
index dae8fe9..44f6d5a 100644
--- a/test/Translation/Forth/exit.forth
+++ b/test/Translation/Forth/exit.forth
@@ -9,5 +9,6 @@
 \ CHECK: ^[[RET]](%[[R:.*]]: !forth.stack):
 \ CHECK:   return %[[R]] : !forth.stack
 
+\! kernel main
 : EARLY-EXIT 1 IF EXIT THEN 42 ;
 EARLY-EXIT
diff --git a/test/Translation/Forth/gpu-indexing.forth b/test/Translation/Forth/gpu-indexing.forth
index fb9bc6f..fe4c289 100644
--- a/test/Translation/Forth/gpu-indexing.forth
+++ b/test/Translation/Forth/gpu-indexing.forth
@@ -13,6 +13,7 @@
 \ CHECK: forth.grid_dim_y
 \ CHECK: forth.grid_dim_z
 \ CHECK: forth.global_id
+\! kernel main
 TID-X TID-Y TID-Z
 BID-X BID-Y BID-Z
 BDIM-X BDIM-Y BDIM-Z
diff --git a/test/Translation/Forth/interleaved-control-flow.forth b/test/Translation/Forth/interleaved-control-flow.forth
index 9e38d5b..a6fe597 100644
--- a/test/Translation/Forth/interleaved-control-flow.forth
+++ b/test/Translation/Forth/interleaved-control-flow.forth
@@ -44,6 +44,7 @@
 \ CHECK-NEXT:  %{{.*}} = forth.drop
 \ CHECK-NEXT:  cf.br ^bb3
 
+\! kernel main
 : multi-while
   BEGIN DUP 10 > WHILE DUP 2 MOD 0= WHILE 1 - REPEAT DROP THEN ;
 
diff --git a/test/Translation/Forth/leave-conditional.forth b/test/Translation/Forth/leave-conditional.forth
index 5d56a65..370c7ba 100644
--- a/test/Translation/Forth/leave-conditional.forth
+++ b/test/Translation/Forth/leave-conditional.forth
@@ -28,6 +28,7 @@
 \ CHECK:     ^bb[[DEAD]](%{{.*}}: !forth.stack):
 \ CHECK:       cf.br ^bb[[JOIN]]
 
+\! kernel main
 10 0 DO
   I 5 = IF LEAVE THEN
   1 DROP
diff --git a/test/Translation/Forth/leave.forth b/test/Translation/Forth/leave.forth
index 242b3c4..528f6fc 100644
--- a/test/Translation/Forth/leave.forth
+++ b/test/Translation/Forth/leave.forth
@@ -12,4 +12,5 @@
 \ CHECK:     ^bb[[EXIT]](%[[B3:.*]]: !forth.stack):
 \ CHECK-NEXT:  return
 
+\! kernel main
 10 0 DO LEAVE LOOP
diff --git a/test/Translation/Forth/loop-index-depth-error.forth b/test/Translation/Forth/loop-index-depth-error.forth
index f8ad22e..fcb1cf9 100644
--- a/test/Translation/Forth/loop-index-depth-error.forth
+++ b/test/Translation/Forth/loop-index-depth-error.forth
@@ -1,3 +1,4 @@
 \ RUN: %not %warpforth-translate --forth-to-mlir %s 2>&1 | %FileCheck %s
 \ CHECK: 'J' requires 2 nested DO/LOOP(s)
+\! kernel main
 10 0 DO J LOOP
diff --git a/test/Translation/Forth/memory-ops.forth b/test/Translation/Forth/memory-ops.forth
index 2167a18..e340df8 100644
--- a/test/Translation/Forth/memory-ops.forth
+++ b/test/Translation/Forth/memory-ops.forth
@@ -9,5 +9,6 @@
 \ Test CELLS produces literal 8 + mul
 \ CHECK: forth.literal %{{.*}} 8
 \ CHECK-NEXT: forth.mul
+\! kernel main
 1 @ 2 3 !
 4 CELLS
diff --git a/test/Translation/Forth/name-mangling.forth b/test/Translation/Forth/name-mangling.forth
index fa00f46..02d7041 100644
--- a/test/Translation/Forth/name-mangling.forth
+++ b/test/Translation/Forth/name-mangling.forth
@@ -9,6 +9,7 @@
 \ Leading digit gets underscore prefix
 \ CHECK: func.func private @_2START
 
+\! kernel main
 : MY-WORD 1 ;
 : UNDER_SCORE 2 ;
 : 2START 3 ;
diff --git a/test/Translation/Forth/nested-control-flow.forth b/test/Translation/Forth/nested-control-flow.forth
index 8cb108c..730b372 100644
--- a/test/Translation/Forth/nested-control-flow.forth
+++ b/test/Translation/Forth/nested-control-flow.forth
@@ -9,6 +9,7 @@
 \ CHECK-NEXT:  %[[L2:.*]] = forth.literal %[[B1]] 2 : !forth.stack -> !forth.stack
 \ CHECK-NEXT:  %[[PF2:.*]], %[[FL2:.*]] = forth.pop_flag %[[L2]] : !forth.stack -> !forth.stack, i1
 \ CHECK-NEXT:  cf.cond_br %[[FL2]], ^bb3(%[[PF2]] : !forth.stack), ^bb4(%[[PF2]] : !forth.stack)
+\! kernel main
 1 IF 2 IF 3 THEN THEN
 
 \ === IF inside DO ===
diff --git a/test/Translation/Forth/param-declarations.forth b/test/Translation/Forth/param-declarations.forth
index 63cdc6c..c15fbd7 100644
--- a/test/Translation/Forth/param-declarations.forth
+++ b/test/Translation/Forth/param-declarations.forth
@@ -4,6 +4,7 @@
 \ CHECK: func.func private @main(%arg0: memref<256xi64> {forth.param_name = "DATA"}, %arg1: memref<128xi64> {forth.param_name = "WEIGHTS"})
 \ CHECK: forth.param_ref %{{.*}} "DATA"
 \ CHECK: forth.param_ref %{{.*}} "WEIGHTS"
-PARAM DATA 256
-PARAM WEIGHTS 128
+\! kernel main
+\! param DATA i64[256]
+\! param WEIGHTS i64[128]
 DATA WEIGHTS
diff --git a/test/Translation/Forth/param-ref-in-word-error.forth b/test/Translation/Forth/param-ref-in-word-error.forth
index f04393c..283e234 100644
--- a/test/Translation/Forth/param-ref-in-word-error.forth
+++ b/test/Translation/Forth/param-ref-in-word-error.forth
@@ -1,5 +1,6 @@
 \ RUN: %not %warpforth-translate --forth-to-mlir %s 2>&1 | %FileCheck %s
 \ CHECK: parameter 'DATA' cannot be referenced inside a word definition
-PARAM DATA 256
+\! kernel main
+\! param DATA i64[256]
 : BAD-WORD DATA @ ;
 BAD-WORD
diff --git a/test/Translation/Forth/plus-loop-negative.forth b/test/Translation/Forth/plus-loop-negative.forth
index 233af08..a087ba8 100644
--- a/test/Translation/Forth/plus-loop-negative.forth
+++ b/test/Translation/Forth/plus-loop-negative.forth
@@ -21,4 +21,5 @@
 \ CHECK-NEXT:  cf.cond_br %[[CROSSED]], ^bb2(%[[POP_S]] : !forth.stack), ^bb1(%[[POP_S]] : !forth.stack)
 \ CHECK:     ^bb2(%{{.*}}: !forth.stack):
 \ CHECK-NEXT:  return
+\! kernel main
 0 10 DO -1 +LOOP
diff --git a/test/Translation/Forth/plus-loop-without-do-error.forth b/test/Translation/Forth/plus-loop-without-do-error.forth
index 06bc674..0b2b64c 100644
--- a/test/Translation/Forth/plus-loop-without-do-error.forth
+++ b/test/Translation/Forth/plus-loop-without-do-error.forth
@@ -1,3 +1,4 @@
 \ RUN: %not %warpforth-translate --forth-to-mlir %s 2>&1 | %FileCheck %s
 \ CHECK: +LOOP without matching DO
+\! kernel main
 +LOOP
diff --git a/test/Translation/Forth/plus-loop.forth b/test/Translation/Forth/plus-loop.forth
index ac7f4ac..c31ac49 100644
--- a/test/Translation/Forth/plus-loop.forth
+++ b/test/Translation/Forth/plus-loop.forth
@@ -26,4 +26,5 @@
 \ CHECK-NEXT:  cf.cond_br %[[CROSSED]], ^bb2(%[[POP_S]] : !forth.stack), ^bb1(%[[POP_S]] : !forth.stack)
 \ CHECK:     ^bb2(%[[B2:.*]]: !forth.stack):
 \ CHECK-NEXT:  return
+\! kernel main
 10 0 DO 2 +LOOP
diff --git a/test/Translation/Forth/scalar-param.forth b/test/Translation/Forth/scalar-param.forth
new file mode 100644
index 0000000..ec99aca
--- /dev/null
+++ b/test/Translation/Forth/scalar-param.forth
@@ -0,0 +1,8 @@
+\ RUN: %warpforth-translate --forth-to-mlir %s | %FileCheck %s
+
+\ Verify scalar param uses i64 argument type.
+\ CHECK: func.func private @main(%arg0: i64 {forth.param_name = "SCALE"})
+\ CHECK: forth.param_ref %{{.*}} "SCALE"
+\! kernel main
+\! param SCALE i64
+SCALE
diff --git a/test/Translation/Forth/stack-ops.forth b/test/Translation/Forth/stack-ops.forth
index 01cccf6..0ce8fca 100644
--- a/test/Translation/Forth/stack-ops.forth
+++ b/test/Translation/Forth/stack-ops.forth
@@ -10,4 +10,5 @@
 \ CHECK: forth.tuck %{{.*}} : !forth.stack -> !forth.stack
 \ CHECK: forth.pick %{{.*}} : !forth.stack -> !forth.stack
 \ CHECK: forth.roll %{{.*}} : !forth.stack -> !forth.stack
+\! kernel main
 1 DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL
diff --git a/test/Translation/Forth/unloop-exit.forth b/test/Translation/Forth/unloop-exit.forth
index e5058a4..c492fa6 100644
--- a/test/Translation/Forth/unloop-exit.forth
+++ b/test/Translation/Forth/unloop-exit.forth
@@ -15,5 +15,6 @@
 \ CHECK: ^bb[[#RET]](%[[R:.*]]: !forth.stack):
 \ CHECK: return %[[R]] : !forth.stack
 
+\! kernel main
 : FIND-FIVE  10 0 DO I 5 = IF UNLOOP EXIT THEN LOOP 0 ;
 FIND-FIVE
diff --git a/test/Translation/Forth/unloop-without-do-error.forth b/test/Translation/Forth/unloop-without-do-error.forth
index b14af8b..bd89d05 100644
--- a/test/Translation/Forth/unloop-without-do-error.forth
+++ b/test/Translation/Forth/unloop-without-do-error.forth
@@ -1,3 +1,4 @@
 \ RUN: %not %warpforth-translate --forth-to-mlir %s 2>&1 | %FileCheck %s
 \ CHECK: UNLOOP without matching DO
+\! kernel main
 UNLOOP
diff --git a/test/Translation/Forth/word-definitions.forth b/test/Translation/Forth/word-definitions.forth
index 6f29e63..590940c 100644
--- a/test/Translation/Forth/word-definitions.forth
+++ b/test/Translation/Forth/word-definitions.forth
@@ -7,5 +7,6 @@
 \ CHECK: }
 \ CHECK: func.func private @main()
 \ CHECK:   call @DOUBLE(%{{.*}}) : (!forth.stack) -> !forth.stack
+\! kernel main
 : DOUBLE DUP + ;
 5 DOUBLE
diff --git a/tools/warpforth-runner/warpforth-runner.cpp b/tools/warpforth-runner/warpforth-runner.cpp
index 7f9d28a..719175e 100644
--- a/tools/warpforth-runner/warpforth-runner.cpp
+++ b/tools/warpforth-runner/warpforth-runner.cpp
@@ -6,7 +6,8 @@
 ///
 /// Usage:
 ///   warpforth-runner kernel.ptx --param 1,2,3 --param 0,0,0,0 \
-///       --grid 4,1,1 --block 64,1,1 --output-param 0 --output-count 3
+///       --grid 4,1,1 --block 64,1,1 --kernel main \
+///       --output-param 0 --output-count 3
 
 #include <cuda.h>
 
@@ -71,6 +72,7 @@ static std::string readFile(const char *path) {
 
 int main(int argc, char **argv) {
   const char *ptxFile = nullptr;
+  const char *kernelName = nullptr;
   std::vector<Param> params;
   Dims grid, block;
   int outputParam = 0;
@@ -108,6 +110,12 @@ int main(int argc, char **argv) {
         return 1;
       }
       outputCount = atoi(argv[i]);
+    } else if (strcmp(argv[i], "--kernel") == 0) {
+      if (++i >= argc) {
+        fprintf(stderr, "Error: --kernel requires a value\n");
+        return 1;
+      }
+      kernelName = argv[i];
     } else if (argv[i][0] == '-') {
       fprintf(stderr, "Error: unknown option %s\n", argv[i]);
       return 1;
@@ -117,12 +125,17 @@ int main(int argc, char **argv) {
   }
 
   if (!ptxFile) {
-    fprintf(stderr, "Usage: warpforth-runner kernel.ptx [--param V,...] "
-                    "[--grid X,Y,Z] [--block X,Y,Z] "
+    fprintf(stderr, "Usage: warpforth-runner kernel.ptx --kernel NAME "
+                    "[--param V,...] [--grid X,Y,Z] [--block X,Y,Z] "
                     "[--output-param N] [--output-count N]\n");
     return 1;
   }
 
+  if (!kernelName) {
+    fprintf(stderr, "Error: --kernel NAME is required\n");
+    return 1;
+  }
+
   if (params.empty()) {
     fprintf(stderr, "Error: at least one --param is required\n");
     return 1;
@@ -151,7 +164,7 @@ int main(int argc, char **argv) {
   CHECK_CU(cuModuleLoadData(&module, ptx.c_str()));
 
   CUfunction func;
-  CHECK_CU(cuModuleGetFunction(&func, module, "main"));
+  CHECK_CU(cuModuleGetFunction(&func, module, kernelName));
 
   // Allocate device buffers and copy data
   std::vector<CUdeviceptr> devicePtrs(params.size());

From 9161ce22f90679d04bff40e12c18d77b65387f96 Mon Sep 17 00:00:00 2001
From: Alex Cameron <asc@tetsuo.sh>
Date: Thu, 19 Feb 2026 20:58:35 +0900
Subject: [PATCH 2/2] refactor: simplify header parsing and add negative tests

---
 gpu_test/conftest.py                          | 31 +++++-----
 lib/Translation/ForthToMLIR/ForthToMLIR.cpp   | 56 +++++--------------
 lib/Translation/ForthToMLIR/ForthToMLIR.h     |  1 +
 .../header-directive-after-code-error.forth   |  6 ++
 .../Forth/header-duplicate-kernel-error.forth |  4 ++
 .../Forth/header-duplicate-param-error.forth  |  5 ++
 .../Forth/header-missing-kernel-error.forth   |  3 +
 .../header-unknown-directive-error.forth      |  4 ++
 8 files changed, 54 insertions(+), 56 deletions(-)
 create mode 100644 test/Translation/Forth/header-directive-after-code-error.forth
 create mode 100644 test/Translation/Forth/header-duplicate-kernel-error.forth
 create mode 100644 test/Translation/Forth/header-duplicate-param-error.forth
 create mode 100644 test/Translation/Forth/header-missing-kernel-error.forth
 create mode 100644 test/Translation/Forth/header-unknown-directive-error.forth

diff --git a/gpu_test/conftest.py b/gpu_test/conftest.py
index 848475e..87fd95e 100644
--- a/gpu_test/conftest.py
+++ b/gpu_test/conftest.py
@@ -284,8 +284,11 @@ def _parse_array_type(type_spec: str) -> int:
     return int(size_str)
 
 
-def _parse_kernel_name(forth_source: str) -> str:
-    """Parse '\\! kernel <name>' from Forth source header."""
+def _iter_header_directives(forth_source: str) -> Generator[tuple[str, list[str]]]:
+    """Yield (keyword, parts) for each \\! directive in the Forth header.
+
+    Strips comments (-- ...) and splits on whitespace. keyword is lowercased.
+    """
     for line in forth_source.splitlines():
         stripped = line.strip()
         if not stripped.startswith("\\!"):
@@ -296,7 +299,14 @@ def _parse_kernel_name(forth_source: str) -> str:
         if not directive:
             continue
         parts = directive.split()
-        if parts and parts[0].lower() == "kernel":
+        if parts:
+            yield parts[0].lower(), parts
+
+
+def _parse_kernel_name(forth_source: str) -> str:
+    """Parse '\\! kernel <name>' from Forth source header."""
+    for keyword, parts in _iter_header_directives(forth_source):
+        if keyword == "kernel":
             if len(parts) < 2:
                 msg = "Invalid header line: expected '\\! kernel <name>'"
                 raise ValueError(msg)
@@ -312,17 +322,8 @@ def _parse_param_declarations(forth_source: str) -> list[tuple[str, int]]:
     Scalar params are not supported by the GPU runner.
     """
     decls = []
-    for line in forth_source.splitlines():
-        stripped = line.strip()
-        if not stripped.startswith("\\!"):
-            continue
-        directive = stripped[2:].strip()
-        if "--" in directive:
-            directive = directive.split("--", 1)[0].strip()
-        if not directive:
-            continue
-        parts = directive.split()
-        if not parts or parts[0].lower() != "param":
+    for keyword, parts in _iter_header_directives(forth_source):
+        if keyword != "param":
             continue
         if len(parts) < 3:
             msg = "Invalid header line: expected '\\! param <name> <type>'"
@@ -362,7 +363,7 @@ def run(
         kernel_name = _parse_kernel_name(forth_source)
         decls = _parse_param_declarations(forth_source)
         if not decls:
-            msg = "Forth source has no 'param' declarations"
+            msg = "Forth source has no '\\! param' declarations"
             raise ValueError(msg)
 
         params = params or {}
diff --git a/lib/Translation/ForthToMLIR/ForthToMLIR.cpp b/lib/Translation/ForthToMLIR/ForthToMLIR.cpp
index 952faa0..4440537 100644
--- a/lib/Translation/ForthToMLIR/ForthToMLIR.cpp
+++ b/lib/Translation/ForthToMLIR/ForthToMLIR.cpp
@@ -203,38 +203,22 @@ LogicalResult ForthParser::parseHeader() {
   const char *bufStart = buffer->getBufferStart();
   const char *bufEnd = buffer->getBufferEnd();
 
-  auto ltrim = [](llvm::StringRef s) {
-    while (!s.empty() && (s.front() == ' ' || s.front() == '\t'))
-      s = s.drop_front();
-    return s;
-  };
-  auto rtrim = [](llvm::StringRef s) {
-    while (!s.empty() &&
-           (s.back() == ' ' || s.back() == '\t' || s.back() == '\r'))
-      s = s.drop_back();
-    return s;
-  };
-  auto trim = [&](llvm::StringRef s) { return rtrim(ltrim(s)); };
+  auto trim = [](llvm::StringRef s) { return s.ltrim(" \t").rtrim(" \t\r"); };
 
   auto splitWS = [](llvm::StringRef s) {
     SmallVector<llvm::StringRef> parts;
-    while (!s.empty()) {
-      while (!s.empty() && (s.front() == ' ' || s.front() == '\t'))
-        s = s.drop_front();
-      if (s.empty())
-        break;
-      size_t i = 0;
-      while (i < s.size() && s[i] != ' ' && s[i] != '\t')
-        ++i;
-      parts.push_back(s.substr(0, i));
-      s = s.substr(i);
+    s.split(parts, ' ', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
+    // Re-split on tabs: filter empty parts produced by tab-only tokens.
+    SmallVector<llvm::StringRef> result;
+    for (auto part : parts) {
+      SmallVector<llvm::StringRef> sub;
+      part.split(sub, '\t', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
+      result.append(sub.begin(), sub.end());
     }
-    return parts;
+    return result;
   };
 
   bool headerEnded = false;
-  bool sawKernel = false;
-  bool sawDirective = false;
   headerEndPtr = bufStart;
 
   const char *lineStart = bufStart;
@@ -272,12 +256,12 @@ LogicalResult ForthParser::parseHeader() {
       }
 
       std::string directive = toUpperCase(tokens[0]);
-      if (!sawDirective && directive != "KERNEL") {
+      if (kernelName.empty() && directive != "KERNEL") {
         return emitErrorAt(lineLoc, "\\! kernel must appear first");
       }
 
       if (directive == "KERNEL") {
-        if (sawKernel) {
+        if (!kernelName.empty()) {
           return emitErrorAt(lineLoc, "duplicate \\! kernel directive");
         }
         if (tokens.size() != 2) {
@@ -285,12 +269,7 @@ LogicalResult ForthParser::parseHeader() {
                              "kernel directive expects: \\! kernel <name>");
         }
         kernelName = tokens[1].str();
-        sawKernel = true;
-        sawDirective = true;
       } else if (directive == "PARAM" || directive == "SHARED") {
-        if (!sawKernel) {
-          return emitErrorAt(lineLoc, "\\! kernel must appear first");
-        }
         if (tokens.size() != 3) {
           return emitErrorAt(lineLoc,
                              "param/shared directive expects: \\! param "
@@ -329,13 +308,9 @@ LogicalResult ForthParser::parseHeader() {
           }
           if (sizeStr.empty())
             return emitErrorAt(lineLoc, "array type requires a size");
-          for (char c : sizeStr) {
-            if (!std::isdigit(static_cast<unsigned char>(c)))
-              return emitErrorAt(lineLoc, "array size must be an integer");
-          }
-          size = std::stoll(sizeStr.str());
-          if (size <= 0)
-            return emitErrorAt(lineLoc, "array size must be positive");
+          if (sizeStr.getAsInteger(10, size) || size <= 0)
+            return emitErrorAt(lineLoc,
+                               "array size must be a positive integer");
           isArray = true;
         } else {
           if (toUpperCase(typeToken) != "I64") {
@@ -357,7 +332,6 @@ LogicalResult ForthParser::parseHeader() {
           decl.size = size;
           sharedDecls.push_back(decl);
         }
-        sawDirective = true;
       } else {
         return emitErrorAt(lineLoc, "unknown header directive: " + directive);
       }
@@ -382,7 +356,7 @@ LogicalResult ForthParser::parseHeader() {
       lineStart = lineEnd + 1;
   }
 
-  if (!sawKernel) {
+  if (kernelName.empty()) {
     auto loc = llvm::SMLoc::getFromPointer(bufStart);
     return emitErrorAt(loc, "\\! kernel <name> is required");
   }
diff --git a/lib/Translation/ForthToMLIR/ForthToMLIR.h b/lib/Translation/ForthToMLIR/ForthToMLIR.h
index 46850da..6c24e6d 100644
--- a/lib/Translation/ForthToMLIR/ForthToMLIR.h
+++ b/lib/Translation/ForthToMLIR/ForthToMLIR.h
@@ -25,6 +25,7 @@ struct ParamDecl {
 };
 
 /// A declared shared memory region: `shared <name> <type>`.
+/// TODO: Not yet consumed — scaffolding for shared memory support.
 struct SharedDecl {
   std::string name;
   bool isArray = false;
diff --git a/test/Translation/Forth/header-directive-after-code-error.forth b/test/Translation/Forth/header-directive-after-code-error.forth
new file mode 100644
index 0000000..3236bde
--- /dev/null
+++ b/test/Translation/Forth/header-directive-after-code-error.forth
@@ -0,0 +1,6 @@
+\ RUN: %not %warpforth-translate --forth-to-mlir %s 2>&1 | %FileCheck %s
+\ CHECK: header directive must appear before any code
+\! kernel main
+\! param A i64[4]
+A @
+\! param B i64[4]
diff --git a/test/Translation/Forth/header-duplicate-kernel-error.forth b/test/Translation/Forth/header-duplicate-kernel-error.forth
new file mode 100644
index 0000000..bc4145e
--- /dev/null
+++ b/test/Translation/Forth/header-duplicate-kernel-error.forth
@@ -0,0 +1,4 @@
+\ RUN: %not %warpforth-translate --forth-to-mlir %s 2>&1 | %FileCheck %s
+\ CHECK: duplicate \! kernel directive
+\! kernel main
+\! kernel other
diff --git a/test/Translation/Forth/header-duplicate-param-error.forth b/test/Translation/Forth/header-duplicate-param-error.forth
new file mode 100644
index 0000000..9f6501b
--- /dev/null
+++ b/test/Translation/Forth/header-duplicate-param-error.forth
@@ -0,0 +1,5 @@
+\ RUN: %not %warpforth-translate --forth-to-mlir %s 2>&1 | %FileCheck %s
+\ CHECK: duplicate parameter name: A
+\! kernel main
+\! param A i64[4]
+\! param A i64[8]
diff --git a/test/Translation/Forth/header-missing-kernel-error.forth b/test/Translation/Forth/header-missing-kernel-error.forth
new file mode 100644
index 0000000..e533020
--- /dev/null
+++ b/test/Translation/Forth/header-missing-kernel-error.forth
@@ -0,0 +1,3 @@
+\ RUN: %not %warpforth-translate --forth-to-mlir %s 2>&1 | %FileCheck %s
+\ CHECK: \! kernel <name> is required
+1 2 +
diff --git a/test/Translation/Forth/header-unknown-directive-error.forth b/test/Translation/Forth/header-unknown-directive-error.forth
new file mode 100644
index 0000000..4089597
--- /dev/null
+++ b/test/Translation/Forth/header-unknown-directive-error.forth
@@ -0,0 +1,4 @@
+\ RUN: %not %warpforth-translate --forth-to-mlir %s 2>&1 | %FileCheck %s
+\ CHECK: unknown header directive: BOGUS
+\! kernel main
+\! bogus foo bar