Comment by amkkma | Hacker News Reader

Comment by amkkma | Hacker News Reader

Parent Full thread

Hi Keno,

Thanks for the honest assessment. Do you have any thoughts about correctness/ composability of compiler transforms like AD, reliability of GPU acceleration and predictability of optimizations? (basically what you've discussed in some of your compiler talks).

How is that going to be possible in an imperative language? Right now we have lux.jl, which is a pure by convention DL framework, but that ends up being jax without the TPUs, kernel fusion, branching (Lux relies on generated functions) and copy elision (though this last part is being worked on IIUC).

A bunch of folks in the ML, Probprog and fancy array space have been grappling with things like generated functions, type level programming and such, and were wondering about future directions in this space: https://julialang.zulipchat.com/#narrow/stream/256674-compil... there among other discussions

Edit: re : bandwidth issue Jan Vitek's group is thinking a lot about the verification vs flexibility tradeoff and some people are working on a trait/ static typing system. Maybe something can be done to help them along?

> Thanks for the honest assessment. What about correctness/ composability of compiler transforms like AD, reliability of GPU acceleration and predictability of optimizations? (basically what you've discussed in some of your compiler talks).

I don't think we really have a good answer yet, but it's actively being worked on. That said, I don't think we can be faulted for that one, because I don't think anybody really has a good answer to this particular design problem. There's a lot of new ground being broken, so some experimentation will be required.

> TPUs, kernel fusion, branching (Lux relies on generated functions) and copy elision (though this last part is being worked on IIUC).

We have demonstrated that we can target TPUs. Kernel fusion is a bit of an interesting case, because julia doesn't really use "kernels" in the same way that the big C++ packages do. If you broadcast something, we'll just compile the "fused" kernel on the GPU, no magic required. There is still something remaining, which is that when you're working on the array level, you want to be able to do array-level optimization, which we currently don't really do (though again, the TPU work showed that we could), but is broadly being planned.

> Edit: re : bandwidth issue Jan Vitek's group is thinking a lot about the verification vs flexibility tradeoff and some people are working on a trait/ static typing system. Maybe something can be done to help them along?

We work closely with them of course, so I think there'll be some discussions there, but it's a very tough design problem.

Glad to hear it's being worked on!

> That said, I don't think we can be faulted for that one, because I don't think anybody really has a good answer to this particular design problem.

Agreed! To be clear, If there's any implication of "fault" it was certainly not in a moral sense or even anything around making poor design decisions. Julia's compiler is being asked to do many new things with semantics that necessarily predated many advances in PL.

Re Kernel fusion, there's another piece here, which you may or many not have included in "array-level optimizations". Julia's "just write loops" ethos is awesome, until you get to accelerators...now we're back to an "optimizer defined sub language" as TKF puts it. People like loops and flexibility, Dex, Floops.jl, Tullio, Loopvec and KA.jl show that it's possible to retain structure and emit accelerator-able loopy code. But none of those, except for dex, has a solution for fusing kernels that rely on loops. I'm still using the concept of Kernels, because there's still a bit of a separation between low level CUDA.jl code/these various DSLs and higher level array code, even if not as stark as python or C++.

Would be really cool, if like Dex, there's a plan to fuse these sorts of structured loops as well. Dex does it by having type level indexing and loop effects (they're actually moving to a user defined parallel effect handler system (https://arxiv.org/abs/2110.07493) ...the latter can tell the compiler when it's safe to parallelize and fuse+beta reduce loops. But that relies on structured semantics/effects and a higher level IR than exists in Julia.

Not sure what a Julian solution would look like, if possible. But given the usability wins, it would be great to have in Julia as well.

> But none of those, except for dex, has a solution for fusing kernels that rely on loops.

The LV rewrite will. Some day, I'd like to have it target accelerators, but unlike fusion, I've not actually put any research/engineering into it so can't make any promises.

But my long term goal is that simple loops in -> optimized anything you want out. Enzyme also deserves a shout out for being able to generate reverse mode AD loops with mutation.

to add, as you know, this is part of a more general problem about type level programming vs write your own compiler vs the non composability of DSLs, where Julia folks in various other non ML domains like PPLs and fancy arrays have been wondering about how to do things that get compiled away, without relying on compiler heuristics or generated function blowups: https://julialang.zulipchat.com/#narrow/stream/256674-compil...

Another non ML example I discussed with some Probprog folks is that there was an arxiv review of PPLs and Julian ones that heavily rely on macros don't compose well within and across packages. The same mechanism for composability which Dex uses for parallelism and AD (effect handlers) is what new gen PPLs in jax and Haskell are using for composable transformable semantics, so maybe that's worth looking into.

We've been having some discussions about how to bring that to Julia, but stalled on engineering time and PL knowledge. Eventually wanted to talk to the core team about it with proposal in hand, but never got there. Let me know if you'd like to talk to some of those folks who have been involved in the discussions as you design the new compiler plugin infra.

https://julialang.zulipchat.com/#narrow/stream/256674-compil...

Reply on news.ycombinator.com