Week 6 · Lesson 2

Sugar: ops that are really several ops

A training program needs subtraction, division, exp, relu and matrix multiply on almost every line. None of them is a new idea for your compiler. Week 6 adds them as sugar: one op you can write, defined as a few core ops your compiler already has.

Sugar is a rewrite at parse time

When your parser reads %d = sub %a %b, it can build add a (mul b -1) on the spot and hand that to the rest of the compiler. The type checker, the interpreter, rangeify and the renderer never see a sub. That's the point: nine new spellings, and no new cases anywhere past the parser.

Because the sugar is its expansion, the results are bit-identical to writing the core ops by hand. Hash consing makes it tidier still: two sub %a %b statements expand to the same nodes, so they're one node.

sugar	means	notes
`neg x`	`mul x -1`	i32 or f32. Exact: flipping the sign never rounds.
`sub a b`	`add a (neg b)`	The same as `a - b` in IEEE arithmetic, and it wraps on i32.
`fdiv a b`	`mul a (recip b)`	f32 only. Two roundings, so not always the correctly rounded quotient.
`relu x`	`max x 0`	i32 or f32
`exp x`	`exp2 (mul x 1.442695)`	`log2(e)` as an f32
`log x`	`mul (log2 x) 0.6931472`	`ln 2` as an f32
`floor x`	a 12-line program	f32. Shown below.
`cmpeq a b`	two `cmplt`s and two `where`s	i32 or bool, not f32
`matmul A B`	the week 5 composite	reshape, permute, expand, mul, reduce

fdiv is not division

IEEE division rounds once: 10 / 3 is the f32 nearest 3.333..., which is 3.3333333. fdiv 10.0 3.0 rounds twice, once for recip 3.0 and once for the multiply, and lands one unit higher, on 3.3333335. The spec defines division this way because a reciprocal is cheaper than a divide on most hardware, and tinygrad does the same. It's also why div stays i32-only: week 1 made div on f32 an error, and that test still has to pass.

cmpeq, and why not on f32

Two integers are equal when neither is less than the other: where (cmplt a b) false (where (cmplt b a) false true). For bools, xor (xor a b) true. On f32 that trick breaks. NaN is neither less than nor greater than anything, so the formula would call NaN equal to everything. Rather than pick a rule for that, the spec leaves f32 out. The lab uses cmpeq on integer labels, where it's exactly right.

floor, carefully

floor needs a cast to i32 and back, which truncates toward zero. Below zero, truncation rounds up (-2.5 becomes -2), so a correction subtracts 1. Two edge cases shape the rest. Any f32 with |x| ≥ 2^23 is already a whole number, and casting a big one to i32 would be an error, so those go around the cast. So do inf and NaN.

%x   = param 0 dtype=f32
%nx  = mul %x -1.0
%ax  = max %x %nx
%lt  = cmplt %ax 8388608.0
%big = xor %lt true             # |x| >= 2^23, inf or nan: already an integer
%xc  = where %big 0.0 %x
%ti  = cast %xc dtype=i32       # truncates toward zero
%t   = cast %ti dtype=f32
%up  = cmplt %xc %t             # a negative number was truncated up
%t1  = add %t -1.0
%fl  = where %up %t1 %t
%r   = where %big %x %fl

exp2 and sin use floor internally, so this exact program sits inside both. And it's an example of a function that's defined everywhere and still has a bump: floor -0.0 is +0.0 by this definition, where IEEE's own floor would keep the sign.

Softmax

Slot 0 holds five scores. Store softmax(x) into slot 1, exactly as: e = exp x, s = the sum of e, then fdiv e by s expanded to five elements. The starter stores exp x.

Sugar and gradients

Lessons 6.3 and 6.4 differentiate programs. Sugar costs nothing there either. exp is exp2 of a scaled x, so its gradient is the exp2 rule times the chain rule through the mul, and no "gradient of exp" rule is needed. The same goes for sub, fdiv, relu and matmul. floor's gradient is zero for |x| < 2^23, and that falls out too: its value goes through a cast to i32, and gradients don't flow through integers. (Bigger values go around the cast unchanged, so there the gradient is 1.)

The four math ops are different. They stay as ops until render, because their gradient rules (d/dx 2^x = 2^x ln 2) need to see them whole. Differentiating the polynomial inside exp2 would be both slower and less accurate.

Where to expand

Expand sugar in the parser, right where you build nodes. Keep exp2, log2, sqrt and sin as nodes through type checking and gradients, and expand them when you lower for run or render (or write a C function per op that performs exactly the spec's steps).

Check yourself

What is sub -2147483648 1?

What does the expansion do on i32?

add -2147483648 -1 wraps around to 2147483647, exactly like C's unsigned arithmetic cast back to int32_t.

fdiv 10.0 3.0 gives 3.3333335, and the correctly rounded quotient is 3.3333333. Which is right?

What does sugar mean?

Sugar means its expansion, nothing else. The grader compares bits, and the test sugar_fdiv expects exactly 3.3333335.

Why is cmpeq not defined on f32?

What is cmplt nan 1.0? And cmplt 1.0 nan?

Both are false, so "neither is less" holds, and the formula would say NaN equals 1.0. IEEE says NaN equals nothing, not even itself.

What is floor -2.5?

The cast truncates to -2. Is -2.5 < -2?

Truncation gave -2, which is above -2.5, so the correction subtracts 1: -3.0.

Why does floor skip the cast when |x| ≥ 2^23?

How many mantissa bits does an f32 have?

23. From 2^23 up, the spacing between f32s is 1 or more, so there's no fraction left to drop. Sending those around the cast avoids the out-of-range error for values like 1e10.

Does exp need its own gradient rule?

What does the gradient pass see after parsing?

It only sees core ops and the four math ops. exp2's rule gives 2^y ln 2, and the mul multiplies that by 1.442695 = 1/ln 2, which gives e^x up to rounding.