adityazero
Vx gives you the power to do so, but one can be conservative by putting all the values in the fleet/*.vx files to extreme values. An analogy to that would be when we compile for X86-64, we can specify `-mcpu` otherwise by default the compiler picks a conservative backend.
Vx allows one to define compute elements(number of cores, etc), memory elements (number of distinct memory elements, hierarchy if any) and their relationship (placement, ownership etc). Together they define the toplogoy of a machine. Toplogy goes into machine description files (see https://github.com/vx-lang/Vx/tree/main/fleet) and the compiler uses them to 'monomorphize' the program based on that topology.
What a Pointer Forgets
btw, this is a much better experience than watching Youtube videos for learning.
Even in data center space, I believe several layers can use fixed weights and remaining layers will compensate for the variations. Power savings will be huge so I think there is an incentive to do more of these.
so true. i have attended meetings to decide on meeting topics for the next half.
So I went through the open compiler sources, OpenXLA and the Mosaic dialect inside JAX. They name the parts of the machine they generates code for, and there is useful information in there for anyone writing TPU kernels.
`tpu.log` is `printf`. A tag string plus whatever values you hand it. Next to it sit log_buffer, and trace_start/trace_stop for timing regions. I feel kernel debugging on TPU has a worse reputation than it deserves.
`tpu.weird` takes an `f32` and returns a `bool`. JAX's lowering evaluates `http://lax.is_finite` is its negation. So the TPU has a hardware predicate for "this float is weird."
And the MXU is a FIFO. Push the weights, stream the activations, pop the result. Weight-stationary, more than one per core, 32-bit accumulator enforced by the verifier. "bf16 in, fp32 accumulate" is the only shape the unit offers.
90 operations, eight memory spaces, and a lot of guidance to developers.