{"slug": "a-ptx-compiler-in-python-and-the-kernel-that-hung-without-an-error", "title": "A PTX compiler in Python, and the kernel that hung without an error", "summary": "A developer built a single-pass PTX compiler in pure Python for the Emmy Hackathon 2026, translating Kernel IR JSON into NVIDIA PTX without relying on nvcc or LLVM, and all 15 benchmark problems pass on an RTX 5090. The design avoids the 100–300 ms per-candidate frontend overhead of LLVM-based pipelines like Triton and TVM by walking the IR tree once and appending text, using explicit bit-exact float encoding, cvta.to.global address conversion in the prologue, and a shared ACCUM_OPS table for reductions.", "body_md": "*My experience for the Emmy Hackathon 2026*\n\nThe task looked easy on paper. Take a Kernel IR as JSON, spit out NVIDIA PTX on stdout. One command:\n\n```\npython compile.py path/to/kernel.ir.json > out.ptx\n```\n\nThere was a catch. I had to write the translation myself. Emmy has its own CUDA C++ emitter, and `nvcc` and LLVM are sitting right there, so the shortcut was obvious. It was also banned. Emmy was allowed exactly one job: turning the JSON into typed Python dataclasses with `Graph.from_dict`. Everything after that was on me. Expressions, registers, barriers, all of it. The code lives at [https://github.com/SarthakRawat-1/emmy-hackathon-2026](https://github.com/SarthakRawat-1/emmy-hackathon-2026).\n\nAll 15 benchmark problems pass on a real RTX 5090. I learned a lot getting there, mostly from failures that were small and oddly specific. This post is about those.\n\nTriton and TVM lower through LLVM IR. It's a great pipeline with a real cost: I'd put the frontend overhead at 100 to 300 milliseconds per kernel candidate. Compile once, run a million times, and you'll never notice. An autotuner trying thousands of variants will notice immediately.\n\nSingle-pass lowering in plain Python sidesteps that. The compiler walks the tree once and appends text. It finishes in milliseconds and needs nothing beyond the standard library and Emmy's deserializer. That niche drove most of my design choices.\n\nFour pieces.\n\n`compile.py` is the driver. `main()` deserializes the graph, resolves shapes and lays out shared memory, then calls `assemble_ptx()`. That function writes the header (`.version 8.0`, `.target sm_89`, `.address_size 64`), the shared memory declarations, the kernel signature and the register counts. `emit_prologue()` loads parameters and unpacks thread coordinates. `emit_tile_axes()` handles 2D block scheduling. More on that later.\n\n`compiler/context.py` has `CodegenContext`. It keeps the emitted lines, the map from IR variables to registers (`bind_var`, `get_var`), a label counter and the register banks. `DTYPE_INFO` is a small table from IR scalar types to PTX type strings and byte sizes. The types are `f32`, `s32`, `u32`, `b16`, `f16`, `bf16`, `b64` and `pred`.\n\n`compiler/expr.py` is `lower_expr()`. It covers literals, variables, binary and ternary expressions, casts, flat indexing, builtins and function calls.\n\n`compiler/stmts.py` is `StmtLowerer`, which handles the statements: `Load`, `Write`, `Accum`, `Assign`, `StridedLoop`, `Loop`, `Cond`, `Sync`, `WarpShuffle`, `TreeHalve`, `CpAsyncCopy` and `ZeroPrologue`.\n\nMy favorite bit of the design is `ACCUM_OPS`. It maps each accumulation to a PTX instruction and its identity value:\n\n```\nACCUM_OPS = {\n    \"add\":      (\"add.f32\", 0.0),\n    \"maximum\":  (\"max.f32\", -1e30),\n    \"minimum\":  (\"min.f32\",  1e30),\n    \"multiply\": (\"mul.f32\", 1.0),\n}\n```\n\nSums, softmax row maxes and products all share one path. A new operator costs one line.\n\nKernel buffers show up as `.param .u64`. Leave them alone and you're holding generic addresses, so the hardware has to work out which memory space each one points into on every access. Convert once, in the prologue:\n\n```\nld.param.u64       %rd1, [param_in];\ncvta.to.global.u64 %rd2, %rd1;\n```\n\nEvery buffer gets this before anything else touches it.\n\nI didn't want any parser deciding how to round `0.1`. So `float_to_ptx_hex()` packs the value as a big-endian 32-bit float and prints the hex form PTX accepts:\n\n``` python\ndef float_to_ptx_hex(val):\n    return \"0f\" + struct.pack(\">f\", float(val)).hex()\n```\n\nThe constant in the output is exactly the bits I computed. Integers that meet floats go through `as_f32()`, which emits `cvt.rn.f32.s32` for signed values and `cvt.rn.f32.u32` for unsigned. Mix those up and nothing looks wrong until a value goes negative.\n\nTranscendentals are explicit too. PTX only gives you base-2 approximations. So `exp(x)` becomes a multiply by log2(e), encoded as `0f3fb8aa3b`, then `ex2.approx.f32`. `log` goes through `lg2.approx.f32`, then a multiply by ln(2), `0f3f317218`.\n\nPTX wants a declared count of virtual registers per type, so the context keeps five banks. `%p` for predicates. `%f` for 32-bit floats. `%r` for 32-bit integers and bitfields. `%rd` for 64-bit addresses. `%h` for 16-bit values.\n\nThat last bank exists because of a quirk I only found by getting rejected. Converting `f32` to `bf16` seems like it should write into any register that's wide enough:\n\n```\ncvt.rn.bf16.f32 %h0, %f3;\n```\n\nIt doesn't. Aim it at a `.b32` register and `ptxas` refuses the instruction. The destination must be 16 bits, declared as `.reg .b16 %h<N>;`. So there's an `alloc_b16()` and a separate declaration line. Tiny detail, cost me an evening.\n\nMost of my kernels are pointwise, one thread per element. Grids get rounded up to whole blocks, so a few threads usually have nothing to do. The obvious move is to send them home early:\n\n```\nsetp.ge.s32 %p1, %r_gid, 1024;\n@%p1 bra EXIT;\n```\n\nThat works until the kernel has a barrier. `bar.sync` waits for every thread in the block. If some have already jumped to `EXIT`, the rest wait forever. No crash, no error message. Just a hang.\n\nSo `compile.py` runs `has_collectives()` over the body first. If it finds a `WarpShuffle`, `TreeHalve`, `Sync` or `CpAsyncCopy`, it emits no early exit at all. This has a price, and I'll state it plainly: collective kernels need a launch grid that matches the problem exactly. I took that deal. A compiler that can emit kernels that freeze is worse.\n\nReductions come in two layers. Inside a warp, `WarpShuffle` sweeps deltas of 16, 8, 4, 2 and 1 with butterfly shuffles. Each step swaps a value with a partner lane and runs the combine expression in registers:\n\n```\nshfl.sync.bfly.b32 %r5, %r4, 16, 0x1f, 0xffffffff;\n```\n\nNo shared memory involved. The combine step is whatever AST the IR specifies, so a sum, a max and the paired updates of online softmax all use the same machinery.\n\nAcross warps, `TreeHalve` goes through shared memory. Each round, the stride is `(cur + 1) // 2`. Active threads read their partner's value, run the combine statements and write back, then everyone meets at `bar.sync`. When it's done, thread 0 holds the answer and broadcasts it. That `+ 1` matters. Plain `cur // 2` drops an element when the size is odd.\n\nThe GEMM kernels use `cp.async` to overlap loads with compute. The variant depends on transfer size. Sixteen-byte copies use `.cg`, which goes through L2 and skips L1:\n\n```\ncp.async.cg.shared.global [%r_dst], [%rd_src], 16;\n```\n\nFour- and eight-byte copies use `.ca`, because `.cg` only supports sixteen. Copies are grouped with `cp.async.commit_group;`, and the kernel calls `cp.async.wait_group 0;` before reading anything it asked for. These are collective operations, so they trigger the no-early-exit rule too.\n\nFor matrix multiplication, the order thread blocks run in matters. A plain row-major launch has neighbouring blocks reading very different parts of the operands. `emit_tile_axes()` remaps the linear `%ctaid.x` into 2D super-tiles, using `super_tile_id`, `super_fast_base` and `intra_super_offset`. Adjacent blocks end up working on a contiguous patch of the output at about the same time, so the tiles they share stay warm in L2.\n\nThe output says `.target sm_89`. The runner is an RTX 5090, which is `sm_120`. That's fine, because PTX is an intermediate form. The NVIDIA driver JIT-compiles the `sm_89` text into native code for whatever GPU it finds. I needed nothing architecture-specific beyond that.\n\nThe local suite runs with `python -m unittest discover tests -v`. Sixteen tests cover command-line invocation, register banking, signed versus unsigned promotion, and vector loads and stores. They also check the refusals: non-f32 atomics, float modulo and write swizzling all raise errors instead of producing something wrong. All 16 pass in about 2.6 seconds.\n\nFor hardware, `submit.py` sends a problem to the official runner:\n\n```\npython submit.py <problem_name> --compiler compile.py\n```\n\nAll 15 problems pass.\n\n| Group | Problems | Result | \n|---|---|---|\n| Pointwise | p00_scalar_add, p01_vector_add, p02_scale_shift, p03_relu_mul | Max abs error 0.000e+00 | \n| Reductions | r01_array_sum, r02_big_sum, r03_mean, r04_row_sum | Within 1.5e-05 | \n| Normalization | f01_rmsnorm, f02_row_softmax, f03_l2norm | Within 1.9e-06 | \n| Matmul | m01_matmul_512, m02_matmul_rect | Within 8.3e-05 | \n| Attention | a01_attention, a02_attention_128 | Within 6.5e-07 | \n\nA few gaps, listed so nobody's surprised. `TreeHalve.inner`, the segmented tree reduction, raises `NotImplementedError`. So does `Write.swizzle`. Collective kernels assume the launch grid fits the iteration space exactly. And the ragged-tail logic for swizzled launches is written but untested, since no open benchmark exercises it.\n\nI expected expression lowering to be most of the work. It was the easy part. The pain lived where the hardware has opinions: a register width, a barrier that needs everyone, a copy instruction that only takes certain sizes. A GPU compiler is mostly those opinions, written down in the right order. The repository is linked at the top if you want to see how they landed in code.", "url": "https://wpnews.pro/news/a-ptx-compiler-in-python-and-the-kernel-that-hung-without-an-error", "canonical_source": "https://dev.to/shogun_the_grt/a-ptx-compiler-in-python-and-the-kernel-that-hung-without-an-error-187h", "published_at": "2026-10-28 08:00:44+00:00", "updated_at": "2026-10-28 08:11:33.954508+00:00", "lang": "en", "topics": ["ai-chips", "ai-infrastructure", "developer-tools", "ai-tools"], "entities": ["Emmy Hackathon 2026", "NVIDIA", "PTX", "RTX 5090", "Triton", "TVM", "LLVM", "SarthakRawat-1"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/a-ptx-compiler-in-python-and-the-kernel-that-hung-without-an-error", "markdown": "https://wpnews.pro/news/a-ptx-compiler-in-python-and-the-kernel-that-hung-without-an-error.md", "text": "https://wpnews.pro/news/a-ptx-compiler-in-python-and-the-kernel-that-hung-without-an-error.txt", "jsonld": "https://wpnews.pro/news/a-ptx-compiler-in-python-and-the-kernel-that-hung-without-an-error.jsonld"}}