Skip to content

FlattenCallExpr Pass

Flattens nested call expressions into three-address code form.

Overview

This pass ensures that call expressions do not appear in nested contexts by extracting them into temporary variables. It enforces a three-address code constraint where:

  1. Call arguments cannot be calls
  2. If conditions cannot be calls
  3. For loop ranges (start/stop/step) cannot be calls
  4. Binary/unary expression operands cannot be calls
  5. Return values cannot be calls

Requires: TypeChecked, SSAForm properties. These properties are automatically verified at BASIC level once produced; use a VerificationInstrument via PassContext if you need required properties to be validated before running passes.

When to use: Typically schedule this pass after the type-checking pass and before code generation to simplify downstream analysis and code generation; this ordering is a convention rather than an automatically enforced requirement.

API

C++ Python Level
pass::FlattenCallExpr() passes.flatten_call_expr() Function-level

Factory function:

Pass FlattenCallExpr();

Python usage:

from pypto.pypto_core import passes

flatten_pass = passes.flatten_call_expr()
program_flat = flatten_pass(program)

Algorithm

  1. Detect Nested Calls: Identify call expressions in nested contexts
  2. Extract to Temps: Create temporary variables (named like t__tmp_v0, t__tmp_v1, etc.)
  3. Insert AssignStmt: Add assignment statements before the original statement
  4. Replace with Var: Replace nested call with temporary variable reference
  5. Handle Control Flow: For if/for statements, insert extracted temporaries directly before the control-flow node in the enclosing SeqStmts

Extraction locations:

  • Before AssignStmt/EvalStmt: Insert directly before
  • Before IfStmt/ForStmt: Insert as sibling statements in the enclosing SeqStmts
  • Inside ScopeStmt (pl.at()): Temporaries are always inserted inside the scope body, preserving execution-context boundaries

Example

Nested Call Arguments

Before:

c = foo(bar(a))  # bar(a) is nested in foo's arguments

After:

t__tmp_v0 = bar(a)
c = foo(t__tmp_v0)

Nested Call in If Condition

Before:

if is_valid(compute(x)):
    y = 1

After:

t__tmp_v0 = compute(x)
t__tmp_v1 = is_valid(t__tmp_v0)
if t__tmp_v1:
    y = 1

Multiple Nested Calls

Before:

result = add(mul(a, b), div(c, d))

After:

t__tmp_v0 = mul(a, b)
t__tmp_v1 = div(c, d)
result = add(t__tmp_v0, t__tmp_v1)

Nested in Binary Expression

Before:

x = compute(a) + compute(b)

After:

t__tmp_v0 = compute(a)
t__tmp_v1 = compute(b)
x = t__tmp_v0 + t__tmp_v1

Direct Call in Return

Before:

return self.tile_add(a, b, f)

After:

t__tmp_v0 = self.tile_add(a, b, f)
return t__tmp_v0

Without this normalization, the Call reaches codegen still wrapped inside a ReturnStmt. Some codegen paths (notably OrchestrationCodegen::VisitStmt_(ReturnStmtPtr)) treat ReturnStmt as a no-op and would silently drop the kernel dispatch.

Nested Call inside Scope Block

Before:

with pl.at(level=pl.Level.CORE_GROUP):
    result = assemble(target, cast(x, BF16), offsets)

After:

with pl.at(level=pl.Level.CORE_GROUP):
    t__tmp_v0 = cast(x, BF16)
    result = assemble(target, t__tmp_v0, offsets)

The temporary stays inside the pl.at() block. Without this scope-awareness, t__tmp_v0 = cast(...) would be hoisted outside the scope, causing a "Misplaced tensor op" error during codegen.

Implementation

Header: include/pypto/ir/transforms/passes.h

Pass FlattenCallExpr();

Implementation: src/ir/transforms/flatten_call_expr.cpp

  • Uses IRMutator to traverse expressions
  • Maintains temporary variable counter
  • Collects extracted assignments
  • Rebuilds statements with flattened expressions

Python binding: python/bindings/modules/passes.cpp

passes.def("flatten_call_expr", &pass::FlattenCallExpr, "Flatten nested calls");

Tests: tests/ut/ir/transforms/test_flatten_call_expr_pass.py

  • Tests call arguments extraction
  • Tests if condition extraction
  • Tests for range extraction
  • Tests binary/unary expression extraction
  • Tests multiple nested calls
  • Tests scope-aware extraction (pl.at() blocks)

Error Types

The pass can detect and report nested call violations via NestedCallErrorType:

  • CALL_IN_CALL_ARGS: Call in call arguments
  • CALL_IN_IF_CONDITION: Call in if condition
  • CALL_IN_FOR_RANGE: Call in for range
  • CALL_IN_BINARY_EXPR: Call in binary expression
  • CALL_IN_UNARY_EXPR: Call in unary expression