Polymage Labs ties up with US-based Tenstorrent to develop AI compiler

Polymage Labs, a deep-tech startup incubated at the IISC, has entered into a strategic collaboration with US-based Tenstorrent Inc to develop an AI compiler for the company's accelerators, it said on Monday.

The collaboration aims to combine Polymage Labs' compiler capabilities with Tenstorrent's AI computing platforms to improve software support for advanced AI hardware, it said.

According to Polymage Labs, compiler software plays a key role in simplifying the programming of complex chips by translating high-level developer instructions into optimised code that can run efficiently on hardware.


"Within a few months of collaboration, Polymage Labs' PolyBlocks compiler framework has been integrated with Tenstorrent's hardware, delivering strong performance directly from unmodified, high-level PyTorch and JAX code," the company said in a statement.

It said the partnership addresses a major challenge in the AI hardware industry - the lack of a mature and seamless software ecosystem - which has affected the adoption of several AI accelerator platforms.

"The synergy between the Polymage and Tenstorrent teams has been exceptional," said Uday Bondhugula, Chief Technology Officer of Polymage Labs and Professor of Computer Science at the Indian Institute of Science.

"Using modern MLIR (Multi-Level Intermediate Representation)-based compiler infrastructures in both PolyBlocks and Tenstorrent's tt-mlir (open-source compiler framework) allowed us to develop a complete, end-to-end compiler for PyTorch and JAX in just a few months-an engineering effort that would otherwise take several years when starting from scratch," he said.

The collaboration validates the company's approach and demonstrates how an effective and reusable compiler stack can be built rapidly with the right technologies, building blocks and collaboration, he added.

PyTorch is an open-source deep learning framework originally developed by Meta Platforms and supported by the Linux Foundation.

JAX is a Python-based library designed for high-performance numerical computing and large-scale machine learning, with a focus on accelerator-driven computation and program transformation.