A PTX compiler in Python, and the kernel that hung without an error

My experience for the Emmy Hackathon 2026

The task looked easy on paper. Take a Kernel IR as JSON, spit out NVIDIA PTX on stdout. One command:

python compile.py path/to/kernel.ir.json > out.ptx

There 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.

All 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.

Why bother going direct

Triton 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.

Single-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.

The shape of the compiler

Four pieces.

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.

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.

compiler/expr.py is lower_expr(). It covers literals, variables, binary and ternary expressions, casts, flat indexing, builtins and function calls.

compiler/stmts.py is StmtLowerer, which handles the statements: Load, Write, Accum, Assign, StridedLoop, Loop, Cond, Sync, WarpShuffle, TreeHalve, CpAsyncCopy and ZeroPrologue.

My favorite bit of the design is ACCUM_OPS. It maps each accumulation to a PTX instruction and its identity value:

ACCUM_OPS = {
    "add":      ("add.f32", 0.0),
    "maximum":  ("max.f32", -1e30),
    "minimum":  ("min.f32",  1e30),
    "multiply": ("mul.f32", 1.0),
}

Sums, softmax row maxes and products all share one path. A new operator costs one line.

Small things that bit me first

Pointers must be global

Kernel 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:

ld.param.u64       %rd1, [param_in];
cvta.to.global.u64 %rd2, %rd1;

Every buffer gets this before anything else touches it.

Float literals and the decimal trap

I 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:

def float_to_ptx_hex(val):
    return "0f" + struct.pack(">f", float(val)).hex()

The 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.

Transcendentals 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.

The register that has to be 16 bits wide

PTX 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.

That 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:

cvt.rn.bf16.f32 %h0, %f3;

It 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.

The bug that is really a design rule

Most 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:

setp.ge.s32 %p1, %r_gid, 1024;
@%p1 bra EXIT;

That 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.

So 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.

Three subsystems that took real work

Reductions and online softmax

Reductions 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:

shfl.sync.bfly.b32 %r5, %r4, 16, 0x1f, 0xffffffff;

No 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.

Across 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.

Asynchronous copies

The 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:

cp.async.cg.shared.global [%r_dst], [%rd_src], 16;

Four- 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.

Swizzling blocks for the L2 cache

For 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.

One target, two architectures

The 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.

What I tested

The 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.

For hardware, submit.py sends a problem to the official runner:

python submit.py <problem_name> --compiler compile.py

All 15 problems pass.

Group	Problems	Result
Pointwise	p00_scalar_add, p01_vector_add, p02_scale_shift, p03_relu_mul	Max abs error 0.000e+00
Reductions	r01_array_sum, r02_big_sum, r03_mean, r04_row_sum	Within 1.5e-05
Normalization	f01_rmsnorm, f02_row_softmax, f03_l2norm	Within 1.9e-06
Matmul	m01_matmul_512, m02_matmul_rect	Within 8.3e-05
Attention	a01_attention, a02_attention_128	Within 6.5e-07

What it does not do

A 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.

What I took away

I 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.