Predicting the performance of large-scale distributed machine learning (ML) workloads across multiple accelerator architectures remains a central challenge in ML system design. Existing GPU and TPU focused simulators are typically architecture-specific, while distributed training simulators rely on workload-specific analytical models or costly post-execution traces, limiting portability and cross-platform comparison. This work evaluates whether MLIR’s StableHLO dialect can serve as a unified workload representation for cross-architecture and crossfidelity performance modeling of distributed ML workloads. The study establishes a StableHLO-based simulation methodology that maps a single workload representation onto multiple performance models, spanning analytical, profiling-based, and simulator-driven predictors. Using this methodology, workloads are evaluated across GPUs and TPUs without requiring access to scaled-out physical systems, enabling systematic comparison across modeling fidelities. An empirical evaluation covering distributed GEMM kernels, ResNet, and large language model training workloads demonstrates that StableHLO preserves relative performance trends across architectures and fidelities, while exposing accuracy trade-offs and simulator limitations. Across evaluated scenarios, prediction errors remain within practical bounds for early-stage design exploration, and the methodology reveals fidelity-dependent limitations in existing GPU simulators. These results indicate that StableHLO provides a viable foundation for unified, distributed ML performance modeling across accelerator architectures and simulators, supporting reusable evaluation workflows and crossvalidation throughout the ML system design process.

The rise of Deep Neural Networks (DNNs) has amplified the demand for efficient computation, with General Matrix Multiply (GEMM) operations at their core. While ASICs are efficient but inflexible, GPUs, especially NVIDIA GPUs, equipped with tensor cores, provide a flexible yet high-performance solution for GEMM-based workloads. Previous research and optimizations have largely centered on NVIDIA’s architecture and programming model, which, while effective, can obscure the rationale behind certain design decisions and limit flexibility for further improvements in tensor core designs. In this paper, we present the design and integration of a tensor core into the open-source RISC-V Vortex GPGPU platform, along with a suite of intrinsics designed for GEMM kernel generation. The analysis conducted on the integration elucidates the connections between GPU system architectural parameters and tensor core configuration. We find that the tensor core is severely under-utilized in many cases and that increased compute capacity does not always imply better performance. Hence, we propose a novel technique, cooperative warp execution in tensor core, which leverages hardware-supported warp cooperation within the tensor core to reduce memory requirements for GEMM operations and boost performance over the baseline tensor core implementation by up to 3x.