Representation and JITting

Let's take a closer look at the sphere example from the Ao project page:

(define (sphere r)
    (lambda ( x y z )
        (- (sqrt (+ (* x x) (* y y) (* z z))) r)))

Looking past the blinding walls of parenthesis, this is implementing the function
f(x, y, z) = sqrt(x**2 + y**2 + z**2) - r

Ao implements an f-rep kernel, where shapes are functions of (x, y, z) that return a single value: negative inside the model, positive outside.

This representation dovetails nicely with Scheme as a user-friendly scripting language. For example, the union of two shapes is the minimum of their two functions evaluated at the target point:

(define (union a b)
    (lambda (x y z)
        (min (a x y z) (b x y z))))

Lambda functions are not well-optimized for evaluation in an f-rep kernel, so they're transformed into a more efficient internal representation.

This transformation uses a strategy similar to that of a tracing JIT.

The function is called with special objects that record the sequence of arithmetic operations that they experience, rebuilding the arithmetic tree by tracing the computation.

Let's walk through a basic implementation in pure Scheme.

;; A token is a list with an operation and
;; some number of arguments
(define (make-token opcode . args)
    (cons opcode args))

;; Test to see if the given object is a token
(define (token? t)
    (and (list? t) (symbol? (car t))))

;; wrap-token is a no-op on tokens and
;; wraps numbers in a constant token
(define (wrap-token t)
    (cond ((token? t) t)
          ((number? t) (make-token 'constant t))
          (else (error "Can't wrap" t))))

Here's the most basic implementation of the compiler:

(define (jit f)
    (f (make-token 'x) (make-token 'y) (make-token 'z)))

This evaluation needs to happen in an environment which understands how to combine tokens, so all of the relevant operations are overloaded. Here's an implementation of overloading for addition:

(define +_ +)   ; save the existing addition operator
(define (+ a b) ; then overload it with a token-aware replacement
    (if (or (token? a) (token? b))
        (make-token 'add (wrap-token a) (wrap-token b))
        (+_ a b)))

Testing out this implementation, it correctly converts a lambda function back into its original tree of operations:

scheme@(guile-user)> (define (f x y z) (+ (+ x y) 2))
scheme@(guile-user)> (jit f)
$9 = (add (add (x) (y)) (constant 2))

At first, this seems like a silly thing to do: if we already have all of the original s-expressions, why not convert them directly into the tree? The tracing strategy gives you the rest of Scheme's syntax for free.

Consider the following function:

(define (circle xmin ymin d)
    " Define a circle from its lower-left corner and diameter "
    (let* ((r (/ d 2))
           (x0 (+ xmin r))
           (y0 (+ ymin r))
           (square (lambda (a) (* a a))))
        (lambda (x y z)
            (- (sqrt (+ (square (- x x0))
                        (square (- y y0)))) r))))

When evaluated with this tracing JIT, Scheme handles all of the arithmetic and variable bindings in the let expression; doing it ourselves would involve re-implementing all of that ourselves (with something like the meta-circular evaluator).

The actual JIT is more complicated than the sample code above. In particular, it properly handles n-ary operations (i.e. (+ x y z 2)) and drops identity operations (i.e. (+ x 0)).

The JIT also applies a handful of optimizations when generating the tree:

Back to Ao homepage
Back to blog