IR Parser for DSL Functions

Overview

The IR parser converts Python DSL code to PyPTO IR using decorators (@pl.function, @pl.program). It enforces SSA properties, tracks source locations, and supports nested control flow.

Key components: Decorator → AST Parser → IR Builder → Scope Manager (SSA) → ir.Function

See IR Builder for manual IR construction and Python IR Syntax for full syntax.

Usage

Basic Function

import pypto
import pypto.language as pl

@pl.function
def simple_add(
    x: pl.Tensor[[64, 128], pl.FP16],
    y: pl.Tensor[[64, 128], pl.FP16],
) -> pl.Tensor[[64, 128], pl.FP16]:
    result: pl.Tensor[[64, 128], pl.FP16] = pl.add(x, y)
    return result

# simple_add is now an ir.Function object
assert isinstance(simple_add, pypto.ir.Function)

Type Annotations

All parameters and local variables require type annotations:

x: pl.Tensor[[64, 128], pl.FP16]  # Recommended subscript syntax
x: pl.Tensor((64, 128), pl.FP16)  # Legacy call syntax (also accepted)

Both syntaxes are equivalent; printer always outputs subscript notation.

For Loops with Iteration Arguments

Use pl.range() with tuple unpacking for loop-carried values (iter_args):

for i, (sum_val,) in pl.range(10, init_values=(sum_init,)):
    new_sum: pl.Tensor[[1], pl.INT32] = pl.add(sum_val, i)
    sum_out = pl.yield_(new_sum)  # Use pl.yield_ (not yield)

Syntax: loop_var, (iter_arg1, ...) - number of iter_args must match init_values.

Yielding and If Statements

Use pl.yield_() to return values from nested scopes:

# Single/multiple value yield
result = pl.yield_(expr)
v1, v2, v3 = pl.yield_(expr1, expr2, expr3)

# If statements create phi nodes
if x > 0:
    positive: pl.Tensor[[64], pl.FP32] = pl.mul(x, 2.0)
    result = pl.yield_(positive)
else:
    negative: pl.Tensor[[64], pl.FP32] = pl.mul(x, -1.0)
    result = pl.yield_(negative)

Text-Based Parsing

Parse DSL code from strings or files for dynamic code generation:

Function	Purpose	Example
pl.parse(code)	Parse from string (auto-detects function/program)	result = pl.parse("@pl.function\ndef f(x): ...")
pl.loads(path)	Load from file (auto-detects function/program)	result = pl.loads('kernel.py')

Features:

• Auto-detection: Automatically detects whether code contains @pl.function or @pl.program
• Returns ir.Function or ir.Program based on what's found
• Single function/program per parse (raises ValueError otherwise)
• Produces identical ir.Function/ir.Program objects as decorators
• See examples/utils/parse_from_text.py for examples

Deprecated aliases (still supported):

• pl.parse_program(code) → Use pl.parse(code) instead
• pl.loads_program(path) → Use pl.loads(path) instead

SSA Properties

The parser enforces Static Single Assignment:

Single Assignment: Each variable assigned once per scope

# ✓ Valid
y: pl.Tensor[[64], pl.FP32] = pl.add(x, 1.0)

# ✗ Invalid - SSA violation
y: pl.Tensor[[64], pl.FP32] = pl.add(x, 1.0)
y = pl.mul(x, 2.0)  # Error: y already defined

Scope Isolation: Variables from inner scopes must be yielded

# ✗ Invalid - temp not yielded
for i, (sum_val,) in pl.range(10, init_values=(x,)):
    temp: pl.Tensor[[64], pl.FP32] = pl.add(sum_val, i)
return temp  # Error: temp not in outer scope

# ✓ Valid - explicit yield
for i, (sum_val,) in pl.range(10, init_values=(x,)):
    temp: pl.Tensor[[64], pl.FP32] = pl.add(sum_val, i)
    result = pl.yield_(temp)
return result  # OK

Iteration Arguments: Create new SSA values each iteration via phi nodes.

Span Tracking and Operations

Span Tracking: Preserves source locations for better error messages

• Each IR node includes Span with filename, line/column ranges
• Enables debugging, error reporting, and source-to-IR mapping

Source-map provenance (pl.parse(code, source_map=...)): When code is generated from another source — as @pl.jit does, re-deriving a kernel into a @pl.program string via ast.unparse — a generated_line → (orig_file, orig_line, orig_col) map remaps each span back to the user's real .py. With it, both parse and compile/pass errors point at the original file (e.g. --> kernel.py:8:13) instead of an anonymous <string>, and the error renderer reads the snippet from that real file via linecache. Mapping is statement-granular: a multi-line user statement resolves to its start line, and synthesized statements (loop-bridge assignments, expanded M, N = a.shape, stripped bind_dynamic) have no original line and keep the generated coordinates. @pl.jit builds this map in Specializer.source_map; the default (None) is a no-op for ordinary parsing. See issue #1612.

Supported Operations:

Category	Examples
Tensor Ops	pl.{add, mul, sub, div, matmul, cast, slice, ...}
Binary Expr	a + b, a - b, a * b, a / b, i == 0, x < 10
Literals	42 → ConstInt (INDEX), pl.const(42, pl.INT64) → typed ConstInt, 3.14 → ConstFloat

See Python IR Syntax for complete operation list.

Complete Example

Nested control flow example:

@pl.function
def flash_attn_simplified(
    q: pl.Tensor[[64, 128], pl.FP16],
    k: pl.Tensor[[1024, 128], pl.FP16],
) -> pl.Tensor[[64, 128], pl.FP32]:
    attn_init: pl.Tensor[[64, 128], pl.FP32] = pl.create_tensor([64, 128], dtype=pl.FP32)

    for i, (attn,) in pl.range(16, init_values=(attn_init,)):
        k_block: pl.Tensor[[64, 128], pl.FP16] = pl.slice(k, [64, 128], [i * 64, 0])
        scores: pl.Tensor[[64, 128], pl.FP16] = pl.matmul(q, k_block, b_trans=True)

        if i == 0:
            new_attn: pl.Tensor[[64, 128], pl.FP32] = pl.cast(scores, target_type=pl.FP32)
            result = pl.yield_(new_attn)
        else:
            updated: pl.Tensor[[64, 128], pl.FP32] = pl.add(attn, scores)
            result = pl.yield_(updated)

        final = pl.yield_(result)

    return final

Multi-Function Programs with @pl.program

Define programs containing multiple functions that can call each other:

@pl.program
class MathOps:
    @pl.function
    def square(self, x: pl.Tensor[[1], pl.INT32]) -> pl.Tensor[[1], pl.INT32]:
        result: pl.Tensor[[1], pl.INT32] = pl.mul(x, x)
        return result

    @pl.function
    def sum_of_squares(self, a: pl.Tensor[[1], pl.INT32], b: pl.Tensor[[1], pl.INT32]) -> pl.Tensor[[1], pl.INT32]:
        a_squared: pl.Tensor[[1], pl.INT32] = self.square(a)  # Cross-function call
        b_squared: pl.Tensor[[1], pl.INT32] = self.square(b)
        result: pl.Tensor[[1], pl.INT32] = pl.add(a_squared, b_squared)
        return result

Key Rules:

• Class-based syntax with @pl.program
• Methods require self parameter (automatically stripped from IR)
• Cross-function calls use self.method_name() → resolved to GlobalVar references
• Two-pass parsing: collect GlobalVars, then parse bodies (supports forward references)
• Access functions: program.get_function("name")
• Text parsing: pl.parse(code), pl.loads(path) (auto-detects program/function)
• Printing: program.as_python() generates valid @pl.program class; pass concise=True to omit intermediate type annotations

Examples: See examples/utils/cross_function_calls.py

Limitations and Testing

Current Limitations:

• Scalar comparisons only in if conditions (not tensors)
• No nested function definitions inside @pl.function
• Limited Python subset (no classes, decorators within functions)
• Explicit yields required for all scope outputs
• Type annotations required for all variables

Testing: Run pytest tests/ut/language/parser/ for comprehensive parser tests.

See Also

• Python IR Syntax - Full syntax specification
• IR Builder - Manual IR construction API
• IR Overview - Core IR concepts