Untitled topic
1–2 of 2 posts
Re: undefined
#2So 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.