Searcharxiv⌕ Search
Home/Archive/arXiv

arXiv · 2609.34657

Torch-PIM: Automated Profile-Guided PIM Offloading for PyTorch

Heeeon Lee·Hyunwoo Nam·Junyong Heo·Hyunmo Sung·Jay Hwan Lee·Yeonsoo Kim·Seongho Jeong·Shinhyung Yang·Bernd Burgstaller
View original source ↗Read source document ↗Export metadata ↓

Abstract

Modern deep learning (DL) workloads are limited by data movement, and processing-in-memory (PIM) targets this bottleneck by placing compute units near the memory. However, PyTorch and other DL frameworks lack compiler support for making this decision on the code they lower: existing offloading frameworks target hand-written C/C++ programs, while those that address DL fix the candidate set to a list of operator types before lowering. We present Torch-PIM, a compiler framework that uses profile-guided optimization (PGO) to decide host-versus-PIM placement over the loop nests that progressive lowering materializes. Every parallel loop nest the pipeline emits enters the candidate space, and each is assessed in two stages: the amount of work it carries, and its memory boundedness. Every quantity the assessment consumes is profiled on the host or obtained from the multi-level intermediate representation (MLIR) of the code. Across PIM configurations of 32 to 128 cores, Torch-PIM's offloading decisions yield speedups of up to 8.6x on tensor operators, 2.9x on MLP, 4.4x on Attention, 5.1x on GPT-J-6B, and 3.6x on LLaMA-7B over CPU-only execution.

Explore related subjects

cs.ARcs.DC

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Heeeon Lee, Hyunwoo Nam, Junyong Heo, Hyunmo Sung, Jay Hwan Lee, Yeonsoo Kim, Seongho Jeong, Shinhyung Yang, Bernd Burgstaller. 2026-09-28. Torch-PIM: Automated Profile-Guided PIM Offloading for PyTorch. https://arxiv.org/abs/2609.34657

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

Edge GPU Aware Multiple AI Model Pipeline for Accelerated MRI Reconstruction and Analysis

Advancements in AI have greatly enhanced the medical imaging process, making it quicker to diagnose patients. However, very few have investigated the optimization of a multi-model system with hardware acceleration. As specialized edge devices emerge, the efficient use of their accelerators is becoming increasingly crucial. This paper proposes a hardware-accelerated method for simultaneous reconstruction and diagnosis of \ac{MRI} from \ac{CT} images. Real-time performance of achieving a throughput of nearly 150 frames per second was achieved by leveraging hardware engines available in modern NVIDIA edge GPU, along with scheduling techniques. This includes the GPU and the \ac{DLA} available in both Jetson AGX Xavier and Jetson AGX Orin, which were considered in this paper. The hardware allocation of different layers of the multiple AI models was done in such a way that the ideal time between the hardware engines is reduced. In addition, the AI models corresponding to the \ac{GAN} model were fine-tuned in such a way that no fallback execution into the GPU engine is required without compromising accuracy. Indeed, the accuracy corresponding to the fine-tuned edge GPU-aware AI models exhibited an accuracy enhancement of 5\%. A further hardware allocation of two fine-tuned GPU-aware GAN models proves they can double the performance over the original model, leveraging adequate partitioning on the NVIDIA Jetson AGX Xavier and Orin devices. The results prove the effectiveness of employing hardware-aware models in parallel for medical image analysis and diagnosis.

cs.AR↗
arXiv2026

Bit-Accurate Modeling of GPU Matrix Multiply-Accumulate Units: Demystifying Numerical Discrepancy and Accuracy

Modern AI accelerators rely on matrix multiply-accumulate units (MMAUs), such as NVIDIA Tensor Cores and AMD Matrix Cores, to accelerate deep neural network workloads. MMAUs expose only instruction-level or API-level interfaces of matrix multiply-accumulate (MMA) operations, while leaving internal floating-point arithmetic behavior undocumented. Consequently, MMAUs across vendors and architectural generations often produce numerical discrepancies for identical inputs, and sometimes exhibit reduced numerical accuracy that can cause training instability. Diagnosing and understanding the root causes of these effects is challenging without white-box models of their arithmetic behavior. This paper proposes closed-loop feature probing (CLFP), a generic and systematic framework for constructing bit-accurate arithmetic behavior models of MMA operations. Based on this framework, we analyze all MMA instructions on ten GPU architectures spanning NVIDIA Volta through RTX Blackwell and AMD CDNA1 through CDNA3, and derive the first bit-accurate arithmetic models for these MMAUs. Our models explain previously observed cross-platform numerical discrepancies and accuracy issues, enable white-box numerical error analysis, reveal four types of precision bottlenecks and one type of numerical asymmetry, and inform software workarounds as well as design suggestions for future MMAUs. This work is open-source at https://github.com/microsoft/MMA-Sim

cs.AR↗
arXiv2026

Coarse-to-Fine Macro Placement via Evolutionary Search and Critical Macro Tuning

Macro placement is a critical stage in chip physical design that substantially affects downstream implementation quality. Recent search-based methods improve existing layouts through partial reconstruction, but quality-biased or spatially restricted macro selection can limit the diversity of reconstruction proposals, potentially hindering escape from local optima. Moreover, coarse-grid representations restrict placement precision. To address these challenges, we propose C2FPlace, a \textbf{C}oarse-to-\textbf{F}ine macro \textbf{Place}ment framework that integrates population-based evolutionary search with fine-grained refinement. During coarse-grained optimization, tournament selection chooses promising parents from randomly sampled groups of layouts, and stochastic partial rip-up and re-place generates offspring by sampling macro subsets across the entire layout. A two-phase schedule samples reconstruction ratios from a higher range early in the search and a lower range later, supporting broad exploration followed by more conservative refinement. During fine-grained optimization, critical macro tuning enables positional adjustments beyond the coarse grid to obtain additional half-perimeter wirelength (HPWL) reduction. Experiments on the ISPD2005 benchmark show that C2FPlace reduces HPWL by 17.82\% over EGPlace and 17.86\% over RollPlace on average. On the ICCAD2025 benchmark, C2FPlace achieves the best average ranking among the compared methods under the evaluated power, performance, and area (PPA) metrics. Our codes are available in \href{https://github.com/lxxxxb/C2FPlace}{https://github.com/lxxxxb/C2FPlace}.

cs.AR↗