This session presents the design of a training-aware compilation flow built using MLIR-based infrastructure for lowering PyTorch models toward accelerator-oriented intermediate representations. While many compilation pipelines focus primarily on inference, enabling full training requires handling both forward execution and backward gradient computation within the same framework. The talk explains how models are traced to obtain computation graphs, how forward and backward passes are compiled within a unified pipeline, and how parameter updates and memory dependencies are handled across training iterations. The session focuses on operator mapping strategies and practical considerations when compiling forward and backward graphs consistently in an MLIR-based flow.

Speaker: Mriganka Bezbaruah (CDAC), Akshay K (CDAC), Prachi Pandey (CDAC)