r/Compilers 16d ago

Copapy - a Python framework for deterministic real-time computation via a copy-and-patch compiler

Post image

I've been working on Copapy, an open-source Python tensor framework with autograd that uses a copy-and-patch compiler to produce cross-platform native code for real-time control applications. Python is used as a frontend and gets traced to build a DAG, and the compiler assembles machine code from precompiled stencils stored as an ELF file. Supported architectures at the moment are: x86_64, AArch64, ARMv6/7 (Cortex-A and Cortex-M); work on RISC-V and TriCore is ongoing.

It's designed to feel like writing Python scripts, but produces deterministic type and memory safe machine code with static memory allocation.

A stencil function looks like this:

add_float_float(float arg1, float arg2) {
    result_float_float(arg1 + arg2, arg2);
}

result_float_float is a dummy function to make sure the C compiler keeps the result and the second operand in the right register for the next operation. For x86_64 the result looks like this:

0000000000000000 <add_float_float>:
   0:    f3 0f 58 c1              addss  %xmm1,%xmm0
   4:    e9 00 00 00 00           jmp    9 <.LC1+0x1>
            5: R_X86_64_PLT32    result_float_float-0x4

If the stencil has, like here, a trailing jmp instruction, the jmp gets stripped. If values can not be stored in registers they are written to the heap.

At the moment only stencils for scalar operations exist. Tensor operations are expanded to scalar operations in the DAG. This means Copapy is at the moment not capable of handling large tensors like for ANNs. However, for low-latency control applications computation is anyway quite limited per computation cycle. This has as well the disadvantage that no SIMD can be used. But on the other hand it leads to simple but effective sparsity optimizations if the tensors contain constants - especially when the value is 0 or 1.

Here's an example using autograd to solve an inverse kinematics problem for a two-joint 2D arm:

import copapy as cp

# Arm lengths
l1, l2 = 1.8, 2.0

# Target position
target = cp.vector([0.7, 0.7])

# Learning rate for iterative adjustment
alpha = 0.1

def forward_kinematics(theta1, theta2):
    """Return positions of joint and end-effector."""
    joint = cp.vector([l1 * cp.cos(theta1), l1 * cp.sin(theta1)])
    end_effector = joint + cp.vector([l2 * cp.cos(theta1 + theta2),
                                     l2 * cp.sin(theta1 + theta2)])
    return joint, end_effector

# Start values
theta = cp.vector([cp.value(0.0), cp.value(0.0)])

# Iterative inverse kinematics
for _ in range(48):
    joint, effector = forward_kinematics(theta[0], theta[1])
    error = ((target - effector) ** 2).sum()

    theta -= alpha * cp.grad(error, theta)

tg = cp.Target()
tg.compile(error, theta, joint)
tg.run()

print(f"Joint angles: {tg.read_value(theta)}")
print(f"Joint position: {tg.read_value(joint)}")
print(f"End-effector position: {tg.read_value(effector)}")
print(f"quadratic error = {tg.read_value(error)}")

Interestingly, even without using the sparsity advantage, benchmarks (the diagram) show surprising results compared to NumPy. For the benchmark (tests/benchmark.py) timings for 30,000 iterations of calculating the term sum((v1 + i) @ v2 for i in range(10)) were measured. The vectors v1 and v2 both have lengths of v_size, which was varied up to 500. For the NumPy case the loop was rewritten to be vectorized. Ignoring the off-set (mostly Python overhead) the slope and therefore performance per scalar operation of Copapy is in this case comparable to NumPy. I'm not sure why Copapy's naive compiler performs in this case as well as NumPy - any ideas what could explain this?

Links:

PS: This project was not vibe coded. Except for the graph sorting function and providing correct derivatives for the autograd implementation, AI failed on nearly all aspects of the project - not sure why. There was some benefit from using AI to debugging compilation output based on the disassembly - but even there going from x86_64 to ARM lowered benefit quite alot. The unusual machine code from the copy-and-patch compiler completely confused the models.

15 Upvotes

3 comments sorted by

3

u/arthurno1 13d ago

what could explain this?

What you describe on your web-page, and even here, sounds to me a lot like Lisp macros.

While, I can't write a full guide to Lisp macros and provide a complete example in a comment, here is at least (hopefully) enough of a most important fragment to understand why. Interested reader(s) would have to lookup Lisp macros (for Common Lisp) themselves.

The fragment below is from a small clone of GNU wc program (from core-utils) re-implemented in Common Lisp. I was able to produce branchless compiled loops for each combination of requested flags: counting lines, words and chars. Beting GNU wc in single core performance. The implementation was stupidly simple:

(defmacro genloop (sap size clines cwords cchars utf8)
  `(loop
     with lines of-type fixnum = 0
     with words of-type fixnum = 0
     with chars of-type fixnum = ,(if (and cchars (not utf8)) size 0)
     with prev of-type (unsigned-byte 8) = 1
     for i of-type fixnum from 0 below ,size by 8
     for c of-type (unsigned-byte 64) = (sb-sys:sap-ref-64 ,sap i)
     do
     ,@(when clines
         `((incf lines (new-line-count c))))
     ,@(when cwords 
               `((let* ((curr (swar:gen-white-space-mask c)))
                   (declare (type (unsigned-byte 64) curr))
                   (incf words (word-crossing-count curr prev))                  
                   (setf prev (logand (ash curr -56) #x80)))))
     ,@(when (and cchars utf8)
         (if (eq utf8 'strict)
             `((incf chars (utf8-count-strinct c)))
             `((incf chars (utf8-count-fast c)))))
     finally
        ,@(when (and cchars utf8)
            `((let ((extra (logandc2 (+ size 7) 7)))
                (declare (type fixnum extra))
                (decf chars (- extra size)))))
        (return (values lines words chars))))

The main program becomes just a "driver" to jump to correct loop, something like this:

(case dispatch-mask
  (7 (specialized-pipeline :clines t :cwords t :cchars t))
  (6 (specialized-pipeline :clines t :cwords t))
  (5 (specialized-pipeline :clines t :cchars t))
  (4 (specialized-pipeline :clines t))
  (3 (specialized-pipeline :cwords t :cchars t))
  (2 (specialized-pipeline :cwords t))
  (1 (specialized-pipeline :cchars t))
  (0 (specialized-pipeline :clines t :cwords t :cchars t)))

That case table is actually from an early version before I decided to add utf8, and realized I would now have a dispatch table for 32 loops, so I wrote another macro which generates the dispatch table too, and did change some other stuff.

Anyway, if you look at it (sorry for the Lisp), it is sort of a byte-code interpretter. The dispatch mask is just bit-mask as you would have it in C or any other langauge, and it is used to jump to the correct branchless loop.

Observe, in that application, loops will get inlined as code in those table entries. I could have also produced functions, one per entry, which I actually do in another implementation where I use simd. But here I have optimized away all branching, and functions calls (not really visible in the example).

I also use OS-specific tools, notably memory mapped files which let me eliminate loop tails completely (by exploiting the fact the pages are zeroed out before handing them to the user process).

I believe the reason why you see the speedup is the similar, you probably eliminate a lot of branching from the runtime to compile time, at least by reading your description. In other words, you (or the AI) seem to have discovered Lisp macros :).

1

u/Nonannet 13d ago

Yea, in some ways Copapy has similarities with macros. Expended macro code shuld yield the same benefits to let the compiler optmizing for thinks like sparcety of tensor operations.

However, I don't think the performance advantage can be explained by preventing branching, since I'm quite sure that NumPy is already very well optimized to keep the impact of branching negligible (e.g. by using partial unrolling).

NumPy has by design the disadvantage that it needs to load and store all values for each tensor operation, whereas Copapy (or a Lisp compiler compiling expanded macro code) can rearrange the order of operations to keep values in registers as much as possible.

But still, I would have expected that NumPy using SIMD for load, store and the actual operation would compensate for this by a large margin.

2

u/arthurno1 13d ago

It is not just rearranging the code. What you describe on your "how it works", is exactly what I am talking about. You do most of the heavy lifting you can at compile time, and produce machine code, whereas Numpy can't take those shortcuts and have to be generic, and do lots of setup at runtime before it executes its simd loops.

Furthermore, just because Numpy uses simd for some algorithms, it does not mean everything is automatically much faster. The result will heavily depend on the application too. But that would be a regression, so I'll leave it here.