Enhancement: Support for Tenstorrent Hardware Acceleration via tt-mlir

[p-w-rs]

Tenstorrent hardware offers specialized AI acceleration, and their MLIR-based compiler, tt-mlir, enables targeting this hardware. As far as I can tell Reactant.jl compiles Julia functions into MLIR, including StableHLO (SHLO), presenting an opportunity to accelerate Julia workloads on Tenstorrent devices. Similar to how tt-forge bridges other frameworks like Jax and PyTorch to Tenstorrent hardware, Reactant could extend this capability to Julia.

Development also currently does not require one to buy a tenstorrent card as they now have tensorrent hardware support specifically for developers at koyeb. See the announcement here.

[wsmoses]

That would be awesome! I think we just need to build https://github.com/tenstorrent/tt-xla

want to give things a go?

[p-w-rs]

So I have done a lot of Julia programming (mainly in for deep learning) and I have also done a lot of c and c++ in the past but I have never done anything with compilers, I would be willing to look into this if it sounds feasible for someone who has never heard of SHLO or XLA until today.

[wsmoses]

well I hadn't heard of tt-mlir until today, but happy to try to help if you want to give things an initial go!

I think we essentially need to start with building tt-xla as a jll on yggdrassil (cc @giordano )

[giordano]

I also happen to have access to these boards 🙃 I've been meaning to find the time to play with Reactant on them, but that lacks sadly.

[p-w-rs]

I just signed up with koyeb and requested access but I don't currently have anyway to really get started with testing anything out.

I asked this on the discord:
"I am looking into accelerating Julia code on tenstorrent, there is a package already that can compile Julia code to SHLO and from there to XLA. From what I can tell I could use tt-mlir to turn the SHLO to ttmlir or XLA to ttmlir and be able to run it that way. Does anybody know if this is correct? Would it be better to feed it the SHLO or the XLA?"

And got a response:
"tt-mlir includes passes to convert the SHLO dialect to the TTIR dialect, which can then be lowered to the TTNN dialect. Since SHLO is a high-level dialect, I believe it’s better to feed SHLO."

So I believe the correct way to attempt this is to pass SHLO to ttmlir-opt or ttmlir-translate to then go to TTNN but I am not 100% sure, I think XLA -> TTNN is still experimental.

[giordano]

I have physical cards in my department, so access for me is fairly simple. What I've been meaning to try is to do what Nick Brown did for Fortran via flang: https://www.epcc.ed.ac.uk/whats-happening/articles/xdsl-workshop-hipeac-2025. I only need 72 hours a day to be able to actually work on this.

[giordano]

So I believe the correct way to attempt this is to pass SHLO to ttmlir-opt or ttmlir-translate to then go to TTNN but I am not 100% sure, I think XLA -> TTNN is still experimental.

This approach would be similar to what we used for Julia on IPU in IPUToolkit.jl: there we write the llvm IR to a file and generate the native code using the (clang-based) vendor compiler. The annoying thing is orchestrating the whole program with the provided SDK, which is in C++. I think this approach is doable also for Tenstorrent cards (and equally annoying for the use of C++ SDK) and at least wouldn't require building a new Reactant.

[avik-pal]

So I believe the correct way to attempt this is to pass SHLO to ttmlir-opt or ttmlir-translate to then go to TTNN but I am not 100% sure, I think XLA -> TTNN is still experimental.

This way we will be able to lower to TTNN, but we will have to implement the stack to interact with their compiler.

Using the PJRT client, all we need is to add the API to load the tt plugin (we already do this for TPUs), and then our current stack should work.

[avik-pal]

We should expose the generic non-TPU version of

Reactant.jl/deps/ReactantExtra/API.cpp

Lines 500 to 522 in 52c17c2

	extern "C" PjRtClient *MakeTPUClient(const char *tpu_path, const char **error) {
	// Prefer $TPU_LIBRARY_PATH if set
	std::string tpu_library_path;
	if (auto path = llvm::sys::Process::GetEnv(kEnvTpuLibraryPath)) {
	tpu_library_path = *path;
	} else if (tpu_path) {
	tpu_library_path = std::string(tpu_path);
	} else {
	*error = "Could not find TPU path";
	return nullptr;
	}
	const PJRT_Api *pluginLoad =
	LoadPjrtPlugin("tpu", tpu_library_path.c_str(), error);
	if (pluginLoad == nullptr)
	return nullptr;
	auto tpu_status = InitializePjrtPlugin("tpu", error);
	if (tpu_status)
	return nullptr;
	RegisterProfiler(pluginLoad);
	return GetCApiClient("TPU");
	}

. That way, we can register plugins from julia without having to rebuild reactant

[mofeing]

this should also be applicable to the Metal PJRT plugin

[giordano]

This way we will be able to lower to TTNN, but we will have to implement the stack to interact with their compiler.

Yes, my experience with the IPU stack was that interacting the their compiler and the entire program orchestration wasn't super fun. Side note, I very recently learned that they also use XLA (fun fact: googling for some obscure XLA environment variables brings as first results Graphcore documentation), but I haven't looked into what they expose exactly.

That way, we can register plugins from julia without having to rebuild reactant

If I understand correctly what you're suggesting, that'd be amazing if that means we don't need to manually deal with program orchestration on their side! But it's not clear to me what we'd need to use from their stack in this scenario. And at this point I'm wondering if I could do the same for the IPUs, too!

[avik-pal]

#1241 is a wip version of this. I also added the Metal plugin to that PR as an example of registering new plugins.

[wsmoses]

also intel gpu support would be cool: https://opensource.googleblog.com/2023/06/accelerate-jax-models-on-intel-gpus-via-pjrt.html

[avik-pal]

A lazy way to get rocm & vulkan https://github.com/iree-org/iree/tree/main/integrations/pjrt/python_packages. (we still need a jll of the ireecompiler.so but the plugins are easy to register)