Intermediate representations¶

This page is the reference for Scaly's two intermediate representations (IRs), the expression dialect and the program dialect. The architecture overview explains why there are two and what an IR node is. Each section below lists what a dialect contains and shows how its constructs look when printed.

The expression dialect¶

Expression nodes¶

An Expr has these fields:

sc.format_expr(expr) prints the graph below one node, and sc.render_expr_assembly(fn) prints a Function and every function it calls. The assembly numbers nodes in traversal order, so two runs print the same text.

Types¶

A TensorType holds a fixed shape of non-negative integers, a data type and a diff flag. Shape () is a scalar.

The default is float64. An operation on two data types fails when you build it, and sc.cast converts one explicitly:

x = sc.sym("x", 3)
n = sc.sym("n", 3, dtype=sc.dtypes.int64, diff=False)
x + n                    # TypeError: mixed dtypes float64 and int64; convert one with sc.cast
x + sc.cast(n, "float64")  # float64

sc.cast follows NumPy's astype: a float becomes an integer by dropping its fraction, and a cast to bool is x != 0, so NaN is true. A comparison, sc.isfinite, &, | and ~ give bool. // and % take integers and round the quotient down, as NumPy does, so -7 // 2 is -4 and -7 % 2 is 1. A zero divisor gives 0 for both, as in NumPy. / takes floating operands only. Arithmetic on bool raises and asks for a cast.

A Python number or list beside an expression takes that expression's data type. An int may join an integer or floating expression, but it raises OverflowError outside an integer type's range, and a float beside an integer expression raises TypeError. NumPy arrays and scalars keep their own data type, so they must match the expression. sc.const follows the same rule: it keeps a NumPy value's data type and makes anything else float64.

NumPy's unsigned byte dtype, uint8, is not supported for expressions. A cast, symbol, constant or TensorType requesting that dtype raises ValueError at construction.

n + 1                    # int64
x * np.int64(2)          # TypeError: mixed dtypes float64 and int64
sc.const(np.arange(3))   # an int64 constant
sc.const([0, 1, 2])      # a float64 constant

A numerical call converts each argument to its leaf's data type when NumPy's same_kind casting rule allows it, as from int64 to float64, and raises TypeError for anything else, such as a complex array.

Data types matter inside the graph and the generated procedures. Lowering and numerical evaluation require every Function input and output leaf to be float64. Boolean and integer values inside a function are allowed. The exported C entry reads and writes double buffers, as described in the pointer ABI.

The diff flag says whether a value can carry a derivative. Inputs have it unless created with diff=False. Constants and non-differentiable operations never have it, and any other operation has it when one of its operands does:

x.type.diff                   # True
sc.const(2.0).type.diff       # False
x.floor().type.diff           # True
(x * x.floor()).type.diff     # True

The flag is metadata. The differentiation rules give floor and ceil zero derivative and split ties equally for minimum and maximum. The derivative conventions also cover abs and extremum reductions.

A sparse matrix also carries its stored pattern in its type, and so does a sparse derivative, which is a sparse matrix whose stored values are built from the derivative's coloring.

Operations¶

Arity is the number of operands.

dot, sumsqr, norm_1, norm_2, norm_inf, argmin, argmax, mean, diag and vec are builders that expand into these operations. There are no matrix decompositions or interpolation tables, and no expm1 or log1p.

Calls and mapped calls¶

A symbolic call to a Function adds one call node, and sc.vmap adds one loop node without carries or a condition, whatever the size of the callee. Here a stage model like the one in Getting started is mapped over ten stages and the last state goes through a terminal cost:

@sc.function(sc.arg("z", 2), sc.arg("u", 1), outputs=sc.arg("znext"))
def model(z: sc.Expr, u: sc.Expr) -> sc.Expr:
    return z + 0.1 * sc.concat([z[1:], u])


@sc.function(sc.arg("z", 2), outputs=sc.arg("cost"))
def terminal(z: sc.Expr) -> sc.Expr:
    return sc.sumsqr(z)


@sc.function(sc.arg("zs", 20), sc.arg("us", 10), outputs=sc.arg("J"))
def stages(zs: sc.Expr, us: sc.Expr) -> sc.Expr:
    znexts = sc.vmap(model, 10)(zs.reshape((10, 2)), us.reshape((10, 1))).vec()
    return terminal(znexts[18:])


print(sc.render_expr_assembly(stages))

The module lists model and terminal first. This is the part for stages:

  expr.func @stages(%zs: tensor<20xfloat64 diff>, %us: tensor<10xfloat64 diff>) -> (%J: tensor<float64 diff>) {
    %0 = expr.input {lowering="auto", name="zs"} : tensor<20xfloat64 diff>
    %1 = expr.input {lowering="auto", name="us"} : tensor<10xfloat64 diff>
    %2 = expr.loop(%0, %1) {callee="model", cond=None, in_place=[], index=False, length=10, n_carry=0, output=0, reverse=False, starts=[0, 0], strides=[2, 1]} : tensor<20xfloat64 diff>
    %3 = expr.reshape(%2) {shape=[10, 2]} : tensor<10x2xfloat64 diff>
    %4 = expr.reshape(%3) {shape=[20]} : tensor<20xfloat64 diff>
    %5 = expr.slice(%4) {index=[slice(18, None, None)]} : tensor<2xfloat64 diff>
    %6 = expr.call(%5) {callee="terminal", output=0} : tensor<float64 diff>
    expr.return %6
  }

Both nodes keep the callee as an attribute and select one of its outputs with output. The loop node records the leading-axis slicing as reads of the flat operands. Iteration i reads zs[2i : 2i+2] and us[i : i+1], from starts and strides, and writes two values of the flat output. Outer operands of a mapped loop are always one-dimensional, so the reshapes in stages leave no node behind. The functions guide covers the slicing rules. The program dialect shows what these two nodes become.

Verification¶

Builders reject most invalid nodes as you write them. x.reshape(4) on a three-element x raises a ValueError before any node exists. sc.verify_expr checks a whole graph against the rule table in ir/expr_spec.py, which matters after a pass has rebuilt it. Function construction runs this check automatically, including on graphs rewritten before lowering. To see a failure, build a bad node directly with the Expr constructor, which skips the builder's checks:

x = sc.sym("x", 3)
bad = sc.Expr(sc.ExprOp.RESHAPE, (x,), sc.TensorType((4,)), attrs={"shape": (4,)})
sc.verify_expr((bad * 2.0).sum())
# VerifyError: verify_expr: node %128283723843408 op=reshape failed rule 'reshape-size': RESHAPE source size 3 != target size 4

The message names the node, its operation, the rule and the reason. An unnamed node is labelled with its Python object id, which changes from run to run. Rules that apply to every node, such as the arity check, are in spec_expr_shared. spec_expr adds the rules for each operation.

The program dialect¶

A ProgramNode is a statement, a declaration or a scalar calculation. You do not build programs yourself. Lowering produces them when Scaly generates C, and the listings below show the optimized program for each example.

Program operations¶

Scalar arithmetic reuses the expression dialect's elementwise operations, with two integer divisions. floordiv and floormod round down, as the expression dialect's floordiv and mod do. div and mod on integers round toward zero as C does, and serve index arithmetic, whose operands are never negative.

Calls and loops after lowering¶

Lowering stages from the expression example turns the mapped loop into a program loop around one call and keeps the call:

prog.module {
  prog.proc @model(%z: memref<2xfloat64>, %u: memref<1xfloat64>, %znext: memref<2xfloat64>) {device=host, input_count=2, lowering="auto", scalarize_mode="procedure", scalarized=True, sz_w=0, w_self=0} {
    %v0 = prog.assign prog.load %z[1] : float64 {declare=True}
    prog.store_pair prog.add(prog.load %z[0], prog.mul(0.1, %v0)), prog.add(%v0, prog.mul(0.1, prog.load %u[0])), %znext[0] : float64
  }
  prog.proc @terminal(%z: memref<2xfloat64>, %cost: memref<1xfloat64>) {device=host, input_count=1, lowering="auto", scalarize_mode="procedure", scalarized=True, sz_w=0, w_self=0} {
    %v0 = prog.assign prog.load %z[0] : float64 {declare=True}
    %v1 = prog.assign prog.load %z[1] : float64 {declare=True}
    prog.store prog.add(prog.mul(%v0, %v0), prog.mul(%v1, %v1)), %cost[0] : float64
  }
  prog.proc @stages(%zs: memref<20xfloat64>, %us: memref<10xfloat64>, %J: memref<1xfloat64>) {device=host, input_count=2, lowering="auto", scalarize_mode="disabled", sz_w=0, w_self=0} {
    %s0 = prog.buffer : memref<20xfloat64, private>
    prog.for %it_t0 = 0 to 10 step 1 {kind=global} {
      prog.call @model(%zs[prog.mul(2, %it_t0)], %us[%it_t0], %s0[prog.mul(%it_t0, 2)]) {callee_needs_w=False, n_in=2, n_out=1, w_self=0}
    }
    %t1 = prog.buffer : memref<2xfloat64, private> = prog.alias %s0 offset 18
    %s1 = prog.buffer : memref<1xfloat64, private>
    prog.call @terminal(%t1, %s1) {callee_needs_w=False, n_in=1, n_out=1, w_self=0}
    prog.store prog.load %s1[0], %J[0] : float64
  }
}

The starts and strides of the loop node became the views %zs[2 * it] and %us[it]. A view passed to a call is an address, not a load, and the C for stages shows it as pointer arithmetic (trimmed to the loop and the call):

  double s0[20];
  const double* t1 = s0 + 18;
  double s1[1];
  for (long long it_t0 = 0; it_t0 < 10; ++it_t0) {
    model_raw((arg[0] + (2 * it_t0)), (arg[1] + it_t0), (s0 + (it_t0 * 2)), NULL);
  }
  terminal_raw(t1, s1, NULL);

The two small callees were expanded into scalar statements, which the scalarized=True attribute records. The optimizer may also expand a callee into its caller. If stages returns znexts directly, the loop body contains the arithmetic of model and no call remains. A call node therefore does not promise a separate C function. Lowering and optimization says when each happens.

Reading a lowered program¶

The attributes on a prog.proc line are these:

Views and aliases¶

A view names a location without reading it, and a buffer can alias another buffer at an offset. Together they let a slice or a reshape of a computed value reuse its storage. Here a slice and a reshape feed a callee that takes a matrix:

@sc.function(sc.arg("M", (2, 3)), outputs=sc.arg("s"))
def trace2(M: sc.Expr) -> sc.Expr:
    return M[0, 0] + M[1, 1]


@sc.function(sc.arg("x", 8), outputs=sc.arg("t"))
def f(x: sc.Expr) -> sc.Expr:
    y = x.sin()
    return trace2(y[2:].reshape((2, 3)))

prog.module {
  prog.proc @trace2(%M: memref<2x3xfloat64>, %s: memref<1xfloat64>) {device=host, input_count=1, lowering="auto", scalarize_mode="procedure", scalarized=True, sz_w=0, w_self=0} {
    prog.store prog.add(prog.load %M[0], prog.load %M[4]), %s[0] : float64
  }
  prog.proc @f(%x: memref<8xfloat64>, %t: memref<1xfloat64>) {device=host, input_count=1, lowering="auto", scalarize_mode="disabled", sz_w=0, w_self=0} {
    %s0 = prog.buffer : memref<8xfloat64, private>
    prog.for %i_t0 = 0 to 8 step 1 {kind=global} {
      prog.store prog.sin(prog.load %x[%i_t0]), %s0[%i_t0] : float64
    }
    %t1 = prog.buffer : memref<6xfloat64, private> = prog.alias %s0 offset 2
    %s1 = prog.buffer : memref<1xfloat64, private>
    prog.call @trace2(%t1, %s1) {callee_needs_w=False, n_in=1, n_out=1, w_self=0}
    prog.store prog.load %s1[0], %t[0] : float64
  }
}

%t1 owns no storage. It is %s0 from element 2, and the C is a single pointer (trimmed):

double s0[8];
const double* t1 = s0 + 2;

The reshape costs nothing at all. Buffers are flat and row-major, so the six-element alias is passed as the 2x3 parameter unchanged, and M[1, 1] becomes the flat index 4 inside trace2. Because an alias reads its source's storage, the workspace packer treats a read of %t1 as a read of %s0 when it decides which temporaries can share memory.

Every buffer also has an address space. Parameters are global, which the listing leaves out, temporaries are private, and constant tables are constant.

Loop kinds¶

Each range carries a RangeKind, printed as kind= on a prog.for. The set is modelled on tinygrad's AxisType¹. Lowering and lane widening produce these kinds:

The C renderer emits a global or reduce loop as a plain for. The kind matters to the passes. Loop fusion leaves reduce loops alone, and lane widening splits an independent loop, or a reduction whose order it can keep, into serial chunks of vector lanes. With lanes=4, the loop of x.sin() * x over 64 elements becomes this (trimmed to the loop headers):

prog.for %i_y_start = 0 to 64 step %SCALY_WIDTH_4_4_4_4 {kind=serial} {
  prog.for %i_y_lane = 0 to prog.minimum(%SCALY_WIDTH_4_4_4_4, prog.sub(64, %i_y_start)) step 1 {kind=vector} {

The width is a C macro rather than a number. It is the smaller of the requested lanes and a cap derived from the target's vector registers, so the C preprocessor settles it for each target. Lowering and optimization describes lane widening.

Program verification¶

verify_program applies the same kind of rule table to programs, from ir/program_spec.py. It checks buffer attributes, that views index with scalars, that loads and stores take views, and that loop kinds are valid. A whole module must also have unique procedure names. Lowering verifies the program before and after optimization, so a malformed program fails before any C exists.