Skip to content

From Python to a checked contract

This walkthrough uses the complete, executable VectorAdd showcase. Every Lean excerpt below comes from that file at website build time. The source includes all helper proofs, imports, and audit commands; use the full file when running it. The displayed fragments explain its construction.

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid     = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    # mask  = offsets < n_elements   -- omitted: aligned (n_elements = BLOCK_SIZE)
    x       = tl.load(x_ptr + offsets)
    y       = tl.load(y_ptr + offsets)
    output  = x + y
    tl.store(out_ptr + offsets, output)

The example is an aligned, unmasked program tile. Choose a positive tile size B and require each program’s whole input and output window to be in bounds. Unlike a typical Python vector-add wrapper, it does not mask a partial last tile. The Python decorator and host launch are outside the Lean theorem. For a masked tail, use the separate FlatVectorAdd example.

Floating values in this example are mathematical real numbers. The result does not establish concrete IEEE-754 addition, GPU compilation, or the correctness of a Python host wrapper.

import VeriTile.Triton
import VeriTile.Triton.Memory.Flatten
import VeriTile.Examples.Common
import VeriTile.Meta.StatementAudit

namespace VeriTile.Bench.Examples.VectorAdd

open VeriTile.Triton
open VeriTile.Triton.KernelIO₂ (Implements)
open scoped VeriTile.Triton.KernelIO₂
open VeriTile.Examples
def addKernel (xReg yReg outReg : RegionName) (blockSize : Nat) : ComputeKernel := triton {
  pid  := tl.program_id(0)
  offs := pid * $(blockSize) + tl.arange(0, $(blockSize))
  x    := tl.load($(xReg) + offs)
  y    := tl.load($(yReg) + offs)
  out  := x + y
  tl.store($(outReg) + offs, out)
}

triton { ... } is the embedded DSL. $(blockSize) and $(xReg) interpolate Lean parameters into it. The Python pointer names become named memory regions; pid * B + arange(B) is retained as the program’s address calculation. Lean imports, the namespace, and the parameter types are part of this workflow. This is an explicit translation, not automatic verification of arbitrary Python source.

3. State an independent result and the memory windows

Section titled “3. State an independent result and the memory windows”

The intended value is xs i + ys i, defined from the inputs independently of kernel execution. Wire the translated kernel to its two inputs and output:

def addIO (B : Nat) : KernelIO₂ where
  kernel := addKernel ⟨"x"⟩ ⟨"y"⟩ ⟨"out"⟩ B
  in1 := ⟨"x"⟩
  in2 := ⟨"y"⟩
  out := ⟨"out"⟩
  B := B
  read1 := fun pid => pid * B
  read2 := fun pid => pid * B
  write := fun pid => pid * B

read1, read2, and write give the start of each program’s B-cell window. The KernelIO₂ constructor also requires a proof that algorithm projection succeeds; Lean supplies it for this supported kernel. An unsupported effect cannot enter the contract through an empty-program fallback.

The relation addIO B ⊨ ... quantifies over disjoint flat-memory allocations, program IDs, loaded input windows, and bounds. Under those conditions it requires termination, the specified output values, and preservation of every cell outside the output window. It does not require all other memory or registers to start at zero. B > 0 is an explicit theorem hypothesis.

specification add_kernel_correctness (B : Nat) (hB : 0 < B) :
    addIO B ⊨ fun xs ys i => xs i + ys i := by
  refine KernelIO₂.Implements.intro _ ?_ ?_ ?_
  · exact addKernel_flattenOk ⟨"x"⟩ ⟨"y"⟩ ⟨"out"⟩ B
  · intro bounds s h1 h2 h3
    exact addKernel_traceSafe ⟨"x"⟩ ⟨"y"⟩ ⟨"out"⟩ B bounds s h1 h2 h3
  · intro s₀ xs ys hx hy
    exact addKernel_region_run B hB s₀ xs ys hx hy

The three obligations have different jobs:

Helper What it establishes
addKernel_flattenOk The kernel lies in the fragment supported by the region-to-flat-memory bridge.
addKernel_traceSafe Every load and store in the actual execution stays within the supplied bounds.
addKernel_region_run The region model terminates, stores the pointwise sum, and preserves other cells.

The complete source proves the last obligation from a value lemma and a scatter-store frame lemma. Implements.intro transports those results to the public flat-memory contract. The helper proofs are necessary: the short headline is not a tactic that proves an arbitrary kernel automatically.

Install the pinned Lean toolchain using the setup instructions, then run from the repository root:

Terminal window
lake build
lake env lean bench/examples/VectorAdd.lean

The file includes these gates:

#axiomsClean add_kernel_correctness

/- The headline's statement surface is the IO signature plus the audit-once
Hoare-triple combinator — no other project constant. -/
#stmtSurfaceSubset add_kernel_correctness ⊆
  [addIO, VeriTile.Triton.KernelIO₂.Implements, VeriTile.Triton.KernelIO₂.B]

The first rejects hidden sorry and additional axioms. The second checks the headline’s project constants against the declared surface. Neither check proves that a translation of Python expresses the intended behavior; the scope and memory contract still need review.

For the same comparator gate used by CI, install the tools on Linux following the comparator setup guide. With comparator, lean4export, and landrun on PATH and the systemd user service available, run:

Terminal window
python3 scripts/check_comparator.py --file bench/examples/VectorAdd.lean --trust

Success prints the accepted theorem count and a log directory. Missing tools, failed compilation, a rejected dependency, or failed proof replay make the command fail. This command freezes and replays the current sources. For an agent editing a proof, use scripts/prove.sh to compare against the task snapshot from before the agent’s edits.

The repository’s recorded-demo check creates a temporary copy, changes out := x + y to out := x - y, and leaves the addition contract and proof unchanged. Run it without editing the canonical example:

Terminal window
python3 site/scripts/record-home-demo.py
python3 site/scripts/record-home-demo.py --check

The recorder requires the original proof to succeed and the subtraction mutation to fail with unsolved proof obligations. It writes the diagnostics and source fingerprint to site/src/lib/vector-add-record.json. Separate four-lane evaluations illustrate the difference; those sample values are not the universal correctness proof. The homepage displays this recorded result, not a live browser proof session.

When adapting the example, review the operation, mask, addresses, numeric model, preconditions, and mathematical target together. Keep a failing mutation that would expose the mistake you intend the contract to detect. Use the proof coverage table to inspect the corresponding scope of other ports.