COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
Authors: Chamika Sudusinghe
Gerasimos Gerogiannis
Damitha Lenadora
Charles Block
Josep Torrellas
Charith Mendis
Venue: ICML 2025 poster (2025)
Keywords: machine learning for systems
transfer learning
hardware accelerators
learned cost models
sparse tensor programs
Primary area: applications->everything_else
TL;DR: A novel framework that leverages the homogeneity of input features across hardware platforms while mitigating heterogeneity to efficiently fine-tune learned cost models for accelerating sparse tensor programs on emerging hardware platforms.

Abstract: Sparse tensor programs are essential in deep learning and graph analytics, driving the need for optimized processing. To meet this demand, specialized hardware accelerators are being developed. Optimizing these programs for accelerators is challenging for two reasons: program performance is highly sensitive to variations in sparse inputs, and early-stage accelerators rely on expensive simulators. Therefore, ML-based cost models used for optimizing such programs on general-purpose hardware are often ineffective for early-stage accelerators, as they require large datasets for proper training. To this end, we introduce COGNATE, a novel framework that leverages inexpensive data samples from general-purpose hardware (e.g., CPUs) to train cost models, followed by few-shot fine-tuning on emerging hardware. COGNATE exploits the homogeneity of input features across hardware platforms while effectively mitigating heterogeneity, enabling cost model training with just 5% of the data samples needed by accelerator-specific models to achieve comparable performance. We conduct extensive experiments to demonstrate that COGNATE outperforms existing techniques, achieving average speedups of 1.47× (up to 5.46×) for SpMM and 1.39× (up to 4.22×) for SDDMM.

COGNATE:
Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
Chamika Sudusinghe' Gerasimos Gerogiannis Damitha Lenadora Charles Block
Josep Torrellas'
Charith Mendis"
Abstract
kernels (Kjolstad et al., 2017: Ye et al., 2023; Hong et al., 2019; Jiang et al., 2020). However, the varying non-zero distributions in input sparse matrices have made diffi- cult to develop performance optimizations for sparse tensor programs that consistently work well across diverse inputs.
Sparse tensor programs are essential in deep learn ing and graph analytics, driving the need for op- timized processing. To meet this demand, spe- cialized hardware accelerators are being devel- oped. Optimizing these programs for acceler- ators challenging for two reasons: program performance highly sensitive to variations in sparse inputs, and early-stage accelerators rely expensive simulators. Therefore, ML-based cost models used for optimizing such programs on general-purpose hardware are often ineffec- tive for early-stage accelerators, as they require large datasets for proper training. To this end, we introduce COGNATE, a novel framework that leverages inexpensive data samples from general- purpose hardware (e.g., CPUs) to train cost mod- els, followed by few-shot fine-tuning emerging hardware. COGNATE exploits the homogeneity of input features across hardware platforms while effectively mitigating heterogeneity, enabling cost model training with just 5% of the data samples needed by accelerator-specifie models to achieve comparable performance. We conduct extensive experiments to demonstrate that COGNATE out- performs existing techniques, achieving average speedups of 1.47x (up to 5.46x) for SpMM and 1.39x (up to 4.22x) for SDDMM.
To overcome this challenge, machine learning (ML)-based program optimization techniques have been introduced to optimize sparse tensor programs on established hardware platforms (e.g., CPUs) (Won et al., 2023; Yang et al., 2023). These techniques adaptively select a program configuration based on the input sparse matrix features. For example, WACO (Won al., 2023) introduces learned cost models to predict the runtime cost of programs under different sparse matrices and program configurations. Then, search-based techniques are used to automatically find the optimal pro- gram configuration using the cost model's output. Overall, these ML-based techniques show superior performance and adaptability across a diverse range of inputs compared to manually crafted performance optimization techniques.
Recently, the hardware front, new domain-speciin ma- chines specifically designed for sparse operations are emerg- ing (Pal et al., 2018; Hegde et al., 2019; Aananthakrish- nan et al., 2023; Gerogiannis et al., 2023; Li et al., 2023; Muñoz-Martinez et al., 2023; Jin et al., 2024). These ma- chines, known as hardware accelerators, offer significant speedups over established hardware platforms. Similar to CPUs, sparse accelerators also provide various program configurations (Gerogiannis et al., 2023; Jin et al., 2024; Gerogiannis et al., 2024), which must be configured by soft- ware to achieve the hardware's 's full potential. is important that this potential is tested during early-stage hardware de- velopment (i.e. before the actual chip is available) to inform better hardware design decisions. For example, hardware ar- chitects face the risk of overprovisioning hardware resources (e.g., increasing cache size) to address inefficiencies that could potentially be resolved through improved software strategies (c.g., adopting better tiling strategy). Therefore, it crucial to automatically select the optimal program con- figuration during the design space exploration (DSE) phase of accelerator development.
1. Introduction
Sparse tensor programs have gained increased significance with the recent advancements in sparse deep learning and graph analytics (Beltagy et al., 2020; Ye & Ji, 2021; Child al., 2019; Dao al., 2021) workloads. As a result, many hand-crafted performance optimization techniques have been suggested to improve the performance of sparse
University of Illinois Urbana- Champaai,, USA. Correspon- dence to: Chamika Sudusinghe <chamika2 @ willii@iiiiis..........
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
However, finding the best program configuration for given

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
early-stage hardware accelerator is challenging. This is because software developers have access only to expen- sive simulators and, therefore, cannot utilize ML-based cost model development and auto-tuning techniques commonly used for established hardware platforms, which rely on large supervised datasets. Such datasets often include hundreds of thousands of labeled examples (Won et al., 2023). Unfor- tunately, the time needed collect datasets of similar scale for emerging accelerators relying on simulators is many orders of magnitude longer than the real execution of the program on the actual chip. For example, it can take up to two weeks to collect a single data point using the simulator of the state-of-the-ar SPADE sparse accelerator (Gerogian- nis et al., 2023). At the same time, the same program would take less than a second to execute on the real chip once fabri- cated. Collecting large datasets would require huge clusters running simulations for months even years. Therefore, in order to bring the same benefits of ML--based optimizations to accelerator platforms at their early stages, we need to rethink learned cost model construction techniques that are data- lata-frugal and highly sample-effieient
fective knowledge transfer only between hardware platforms of the same type (e.g. CPU-to-CPU, GPU-to-GPU) (Sasaki et al., 2022; Won et al., 2023; Zheng et al., 2021). Trans- ferring between hardware of different types (e.g. CPU-to- accelerator) poses unique challenges:
Heterogeneous program configuration spaces. The pro- gram configurations for emerging sparse accelerators, which serve as the input feature space for cost models, can dif- fer significantly from those of general-purposs hardware. For example, emerging sparse accelerators have software- managed buffers instead of hardware- managed specialized, rather than mannwade-managed caches and pipelines. This causes a disparity in program configuration spaces for general-purpon hardware and emerging accelerators, mak- ing it challenging to naively apply transfer learning. Exist- ing heterogeneous transfer learning techniques (Liang et al., 2019), such as feature augmentation (Daumé III, 2009; Duan et al., 2012), can be a viable approach. However, these tech- niques often produce feature representations that are too sparse for the cost model to effectively learn, specifically when accommodating diverse set of program configura- tion across different hardware platforms. Figure 2 shows the results of applying popular heterogeneous transfer learning techniques feature augmentation (FA) and feature map- ping (FM) to learned cost model, WACO (Won et al., 2023). Even when using data samples from 1000 matrices for fine-tuning on the SPADE accelerator, the best configura- tions found under WACO+FA and WACO+FM are far from optimal. Therefore, we need better techniques to handle the heterogeneity of program configurations across hardware.
Source Cost Model
Target Cost Model
(Prddicted Best Confi
Sparse Program
Program
Confira
Cortoe
Sparse
Rustines
Perpra
Emerging Hardwars Datiseet
Simulator of Dedicated Sparse Accelerator
Hortware Datasat
Train Source
Target
Matric
Runtime Execution
Figure I: Transfer learning pipeline of COGNATE.
Inspired by the suc-
cess of transfer learn- ing in other domains
High sample efficiency requirement. Existing transfer learning solutions for learned cost models operating in ho- mogeneous feature spaces typically require at least 25% of the original dataset used in a non-transfer learning setup to achieve competitive performance on the target hardware plat- form (Sasaki et al., 2022). The target dataset requirement for these solutions can further increase due to the heterogeneous input feature spaces between general-purpose hardware and emerging accelerators. This makes it infeasible to adopt ex- isting solutions in their current form for accelerators in early design stages. Therefore, we need data- frugal techniques.
1.5
WIKG-FA -Meaie 1.0
Geomean Speedus
--rura Tisp-21
(Weiss et al., 2016;
1.4
Zhuang et al., 2020), re- searchers have proposed
1.3
many techniques to re- duce data requirements for training cost mod- els (Sasaki et al., 2022:
1.2
1.1
1.0
5
100
1000 of Matrices
Zheng et al., 2021). These techniques lever- age knowledge trans- ferred from cost models learned on one hardware platform (source) to an- other (target) using the ubiquitous pre-train and fine-tune paradigm (Krizhevsky et al., 2012). Such techniques have shown to reduce the data requirement for the target platform. Therefore, using such techniques to transfer learn cost models from general- purpose hardware to emerging accelerators can reduce the data requirements from expensive simulations (Figure 1). However, we notice that most prior works have achieved ef-
Performance of existing systems for SpMM on SPADE; WacoNet with feature augmentation (WACO+FA), and feature mapping (WACO+FM).
Figure 2:
COGNATE. In this paper, we present COGNATE, a novel framework for developing learned cost models that enable effective knowledge transfer (Figure 1) overcoming these challenges. COGNATE uses WACO's cost model architec- ture (Won et al., 2023) as the base model ( (WacoNer) but incorporates key changes to make it amenable for transfer learning. This enables the discovery of better program con- figurations that are closer to the optimal (Figure 2), while requiring significantly less fine-tuning data samples.
COGNATE is centered around two key principles intro- duced in (Neyshabur et al., 2020): feature reuse and the
2

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
capture of low-level statistical information. We observe that, even though the program configurations between general- purpose hardware and accelerators are heterogeneous, there are certain feature spaces that can be mapped due to their similarities. Motivated by this observation, we propose an approximate mapping of comparable code optimizations, effectively segregating the feature space generated by pro- gram configuration into homogeneous and heterogeneous components. This allows feature reuse across the source and target platforms. The heterogeneous components repre- sent non-mappable hardware specific parameters that can be disparate across different platforms. Such components can introduce challenges during transfer learning due to negative transfer. To separately encode the heterogeneous feature spaces, we introduce a novel latent space representation of the heterogeneous input feature space using an auto- encoder. This novel formulation of the feature space allows us to effectively reuse features while minimizing the impact of negative transfer. Further, we observe that certain layers of WacoNer do not contribute heavily to the final prediction and this over- parametcriizaioonn can hinder transferability due to wetr-fitting. To mitigate this, COGNATE modifies WacoNet by reducing the number of layers and extracting features at various depths and scales, effectively allowing the model to capture low-level statistical information.
We evaluate and show that COGNATE produces highly accurate transfer learned cost models for emerging sparse accelerators with minimal data collection overhead. Fur- thermore, we perform extensive experiments and ablation studies to demonstrate its benefits and generalizability.
2. Background and Related Work
2.1. Sparse Tensor Programs
Sparse tensor programs perform computational tasks that involve tensors where most of the elements are zero. These computations are optimized to efficiently process only the non-zero values. We describe two operations frequently used in these computations below.
Sparse Matrix-Matrix Multiplication (SpMM) is the op- eration of multiplying a sparse matrix A E IR MxK with a dense matrix B E IR K resulting in an output matrix D E R MxN, The SpMM operation can be expressed as DI,t= = E; Ar. By,s, where Ar, # 0.
Sampled Dense-Dense Matrix Multiplication (SDDMM) is an operation that involves the multiplication of two dense matrices, followed by an elementwise multiplica- tion with a sparse matrix. Given a sparse matrix A E MxN, a sparse output matrix D E IR MxN, and two dense matrices B E R and C E RKXN SD- DMM operation can be expressed as Di,K = Aix E3 (Bug Cz,),, where Ai, A O.
We evaluate COGNATE on two widely used sparse opera- tions, Sparse Matrix-Matre Matrix Multiplication (SpMM and Sampled Dense-Dent Dense Matrix Multiplication (SDDMM). Starting with a CPU as the source hardware platform, we transfer program program configurations the the state-of- the art SPADE (Gerogiannis et al., 2023) emerging sparse accelerator. SPADE's open ISA and vast set of possible pro- gram configurations make it an ideal target platform for our evaluation. Further, to demonstrate the generalizability of COGNATE, we evaluate our techniques for a second target accelerator an NVIDIA A100 GPU.
2.2. Sparse Tensor Programming Systems
Table 1: Program configuration parameters (Config Params) available across CPU, GPU, and SPADE.
Config Params CPU GPU SPADE
Type
Loop Strip-mining
Numerical
Loop Reordering
Categorical
Format Reordering
Categorical
Loop Binding
Categorical
Our results show that COGNATE outperforms existing transfer learning techniques by 28.44%, achieving an aver- age speedup of 1.47x (up to 5.46x) for SpMM and 1.39x (up to 4.22x) for SDDMM on SPADE. On the A100 GPU, it attains an average speedup of 1.17x (up to 1.61x) for SpMM and 1.15x (up to 1.49x) for SDDMM.
Loop Unrolling
Categorical
Tiling
Numerical
Barrier
Binary
Cache Bypassing
Binary
Matrix Reordering
Binary
A sparse tensor programming system supports a range of code optimizations that modify the structure of the code to enhance performance. The effectiveness of these code optimizations depends on assigning specific values to the parameters of the program configuration. By tuning these parameters, we can significantly impact the runtime perfor- mance of sparse operations. Table outlines the parameters available for program configurations across different hard- ware platforms explored in this work (code optimizations are detailed in Appendix B). The execution strategy for sparse tensor programs depends on both the hardware platform and the corresponding programming system used. In this
In summary, this paper makes the following contributions.
We introduce techniques to segregate and encode the ho- mogeneous (approximate mapping of comparable code optimizations) and heterogeneous (latent representa- tion using an auto-encoder) components of program configurations across different hardware platforms.
We introduce COGNATE, novel data-frugal framework for developing learned cost models that are amenable to few-shot fine-tuning across different hardware platforms, leveraging the above techniques.
3

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
work, for CPU execution (source platform), we use TACO (Kjolstad et al., 2017), domain-specifie language and compiler designed for sparse tensor algebra. Considering our target accelerators. SPADE has its own tile-based open instruction set architecture (ISA) to leverage different vari- ations of SpMM and SDDMM operations. For GPU, we employ SparseTIR (Ye et al., 2023), a sparse tensor com- pilation framework developed as an enhancement to TVM Tensor IR (Chen et al., 2018a).
rized into two main types: homogeneous transfer learning (Zhuang al., 2020), where the input and label spaces are the same, and heterogeneous transfer learning (Day & Khoshgoftaar, 2017), where either one or both can be dif- ferent. In program optimization, transfer learning has been successfully used to transfer cost models learned from one hardware platform to another, primarily in homogeneous settings, to minimize the target domain data requirements (Zheng et al., 2021; Ryu & Sung, 2021; Sasaki et al., 2022). In this work, we seek to minimize the target domain data requirement during fine-tuning (Shen et al., 2021), by tar- geting heterogeneous input feature spaces present between general- purprol- hardware and emerging sparse accelerators (data-collection objective) (detailed in Appendix A).
2.3. ML-based cost models
Learned Cost Models. Cost models act as fast cost- effective proxies for executing workloads on real hardware. Their primary goal is to accurately estimate the execution time of workloads as they would perform on real hardware. To achieve this, these cost models can be trained data sam- ples with various program configurations and then be used to predict the program configuration that will deliver the op- timal performance. Hence, generally, the training objective of cost models is tied with minimizing I - t*1, where t" is the runtime of the true optimal program configuration and bear is the runtime of the best program configuration suggested by the cost model (accuracy objective ( (detailed Appendix A). Finding the best configuration suggested by the cost model is usually done using auxiliary intelligent search techniques such as simulated annealing, Monte Carlo tree search, and reinforcement learning. There have been numerous works on learned cost models to predict the run- time of workloads targeting different hardware platforms (Chen et al., 201 8b; Adams et al., 2019). These techniques range from simple XGBoost (Chen& & Guestrin, 2016) based cost models (Chen et al., 201 8b;a) to sophisticated deep neu- ral network based models (Baghdadi et al., 2021; Kaufman et al., 2021; Zhai et al., 2023; Zheng et al., 2020).
3. Our Methodology: COGNATE
Program
Program
(1)
Contipuatioss
Configurations
12)
Homeggenuue
Configuration Mapper
Palsrogunoss
Autooncodar
Courn Dorks
Letent Enceder
Sparae Latret Conti
Suarne
Spresn
Metris
Cannel
(a) WACO
(b) COGNATE
Input Featurizer
Predictor
Runtime
Gost
Figure 3: A comparative overview of the enhanced cost model design in COGNATE (b) alongside WACO's cost model design (a), highlighting key differences.
Here, we present COGNATE, a novel framework to design data-frugal learned cost models to accelerate the execution of sparse tensor programs on emerging hardware. The fol- lowing subsections outline our contributions toward achiev- ing the objectives set forth in Section 2.3; maximizing the accuracy while minimizing the data collection overhead.
WACO's Cost Model. WACO (Figure 3(a)) (Won et al., 2023) introduced a learned cost model specifically built for sparse tensor programs, which utilizes sparsity patterns as raw input data. WACO's cost model employs submanifold sparse convolution networks (SCNN) (Graham & Van der Maaten, 2017) to extract features using an input featurizer. It leverages a neural network- based program embedder to capture the impact of code optimizations on sparse opera- tions by encoding program configurations into embeddings. These program embeddings are merged with the extracted sparsity pattern features produced by the input featurizer. The merged inputs are then processed through multi-layer perceptron predictor to estimate the execution cost.
3.1. Enhancements to Enable Transfer Learning
We build upon WACO considering it our base model by refining its architecture to better facilitate transfer learning across diverse hardware platforms. These enhancements represent contributions that are orthogonal to WACO's orig- inal scope. Our improved cost model design (Figure 3(b)) is structured around four key components: configuration mapper, input featurizer, latent encoder, and predictor. The configuration mapper (Figure 3()) and latent encoder (Figure 3(b)(2) replace the program embedder in WACO, while the input featurizer (Figure 3(b) has been modi- fied to more effectively capture low-level information from sparsity patterns. Both configuration mapper and input fea- turizer remain consistent across hardware platforms, serving as the components that enable efficient knowledge transfer.
Transfer Learning. Transfer learning is a technique that leverages knowledge gained from a task in a source domain to improve the performance of a related task in a target domain, where data collection can often be challenging (Bozinovski, 2020). There have been many successful ex- amples of transfer learning techniques in a wide range of fields (Weiss et al., 2016). Transfer learning can be catego-
Configuration Mapper (FM).. The configuration mapper captures homogeneity across hardware platforms by pro- cessing program configurations (c,) and their parameters to identify similarities in code optimizations across various
4

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
platforms. We designed it to approximately map similar con- figuration parameters across different hardware platforms (described in Section 3.2 and Section E) to a unified feature space. This is achieved by using explicit mapping func- tions. The resulting parameters are subsequently passed through a multi-layer perceptron (MLP) to produce the final configuration vector pj. In this work, we approximate the code optimizations loop strip-mining and loop reordering as Pj- (M)(( at (()),))) ej). using the mapping functions Φ and TT, introduced in Section 3.2.
another, requiring the combination of multiple other code op- timizations to replicate the same impact. For example, loop strip-mining optimization on CPUs can be closely approxi- mated by collectively applying barrier and riling optimiza- tions in SPADE. By mapping the effects of these optimiza- tions using their program configuration parameters, we can expose patterns that facilitate effective knowledge transfer across hardware platforms. In this section, we present our approaches for approximately mapping loop strip-mining, barrier, and tiling optimizations between CPU and SPADE, and loop reordering optimization between CPU and GPU.
Input Featurizer (IFE), Matrices with identical dimen- sions and non-zero elements can exhibit vastly different sparsity patterns, making it difficult to extract meaning- ful features based only on statistical properties. Building on WACO's input featurizer (Won et al., 2023), we mod- ify the network architecture (Figure 3) to more effectively capture low-level information from sparsity patterns. Our network consists of 12 SCNN layers (compared to 14 layers in WACO), arranged in 4 blocks, each containing 3 sparse convolution layers. At the end of each block, we apply max pooling to condense spatial information. We increase the number of channels across blocks up to 256, whereas WACO had them fixed at 32. These additional channels enables our design to capture hierarchical features more effectively throughout the network compared to WACO. For a given sparse matrix M, our inpul featurizer generates a sparse feature vector SM, expressed as M = (M).
Loop strip-mining is a code optimization that decomposes large software loops into smaller segments to optimize com- putations for memory utilization and cache performance. In our context, is applied loops iterating over the in- dices i, j. and & of matrices in SpMM and SDDMM sparse operations (Section 2)1,, where parameters 7, J, and K are used to split these loops into outer and inner segments, resulting in loop nest of six decomposed loops. The re- sulting loop segments are (2,,22:J.11 f2. k1, k2) and their execution order is denoted by w. In SPADE, we approx- imate this using barrier and tiling optimizations. Tiling decomposes matrix into smaller blocks optimize data reuse in the local memory, while barrier controls the ex- ecution order of tiles. For example, enabling barrier opti- mization pauses the tiles scheduled by control processing element until all previous tiles have been completed (Gero- giannis et al., 2023). Similar to strip-mining parameters, the tiling parameters for é, j. and k indices of matrices are represented in SPADE as Peol: Prow, depiit respectively, while barrier is represented by b, where b = 1 if barrier is enabled, and b = 0 otherwise. Intuitively, tiling divides computations into smaller blocks, while barriers control syn- chronization during execution. By enabling and disabling barriers for various tiling configurations, we can dictate the order of computation. This resembles loop strip-mining and reordering in CPUs, where optimizing loop execution enhances performance and cache utilization. We can approx- imately map tiling and barrier parameters to the correspond- ing strip-mining parameters using the mapping function Φ : {Peot Prow, Suplits b} - J, as follows:
Latent Encoder (CE). We handle the heterogeneity of program configurations across hardware platforms using per-target autoencoders that compress the heterogeneous components of the configurations into compact latent rep- resentations (described in Section 3.3). An autoencoder is trained for each target-sparse primitive pair. During both training and inference, the latent encoder LE processes a configuration (cz), transforming it into latent representa- tion 2j- (c,).. ). that encapsulates the unique characteris- ties of the program configuration.
Predictor (P). As the final component of the cost model, the predictor (Figure 3(b)(4) integrates the three feature vectors from the preceding components into a single uni- fied vector, encapsulating all key information about the sparsity pattern and program configuration. This unified vector (sMr) |,]l=z) is passed through a multi-layer percep- tron (MLP) to eventually output a scalar value representing the predicted execution cost, which can be expressed as f,.e, = P(p; lz;).
b(pear proms Suplis b) = (I, J, KW))
I no Deoly J no Prew, K Rs 8ypliii
Ika, 52, b2, 61, 81, ki] if ftb=1 i = 0
w [bz, in, f2, in J1, ki]
Loop reordering is a code optimization that adjusts the exe- cution order (w) of loops to improve cache efficiency and facilitate parallel processing. It is often applied after loop strip-mining. Here, we examine how it can be approximated for both CPU (a1) and GPU (aa). In CPU, loop strip-mining results in six decomposed loops {81, i2, f1, f2, kı, kz)- Similarly, in GPU, loop strip-mining produces six loop seg- ments, but the loop structure differs (6z, iz j. k1, ke ks
3.2. Exploiting Homogeneity: Approximate Mapping of Comparable Code Optimizations
Different hardware platforms often use distinct program- ming systems, leading to variations in how code optimiza- tions are parameterized (Figure 1). Further, an optimization available in one platform may not be directly available on
5

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
WACO+FA la WACO+FM
No Transfer
Zaro-Shat (CPU)
COGNATE ('bp-1) 2)
COGNATE (Tep-5)
Optimal
1.55
Geomean Speedup
1.50
1.25
1.29
1.19
1.25
1:15
1.22
1oc
1.091.0
1:03
1.00
C:99
0.75
C. 6 67
0.59
0.520.55
0.50
0.25
1.07
0:09
0.00
SpMM (SPADE)
SDDMM (SPADE)
SpMM (GPU) Sparse Tensor Operations
SDDMM (GPU)
Figure 4: Geomean speedups of COGNATE and other techniques, normalized to the baseline.
due to architectural differences. We approximate them us- ing 22(.) function that determines the index of a loop seg- ment and a mapping function ma; (21,d2,.... k2 Wa, - {ii, ...,............................ k3, W follows:
complexity. With the hardware specific optimizations now represented in a unified latent feature space, it becomes possible to identify and leverage similarities in their impact on performance during fine-tuning, Finally, this approach facilitates the seamless integration of emerging hardware platforms into COGNATE, support new target hardware platforms by training new au- toencoders and relying on few-shot fine-tuning, eliminating the need to retrain the source model from scratch (detailed in Appendix C). As long as the overall structure of the sparse tensor program remains consistent, COGNATE can quickly adapt by using small number of new performance samples. In contrast, traditional cost model development approaches would require re-evaluating a large number of configurations (Won et al., 2023).
Tay (fi, J1, k1, w2, (ii, 82, J1, hi, Ka, K3, why))
can extend COGNATE to
1y = Suy (kis) = Oay (1e2) +
Tas J, k 22 k3, Was)
11, 3:3 J ki, liz, k3, was)
j = 1, Sag (i))
Sag (i) +1
3.3. Mitigating Heterogeneity: Latent Encoding of Hardware-speciffe Code Optimizations
While use the strategies described in Section 3.2 to approximate code optimizations with homogeneity, such techniques are not applicable to hardware-speciie optimiza- tions. An existing approach for representing hardware- specific optimizations across hardware platforms is to en- code them using feature augmentation. However, this results in excessively sparse feature vectors, as code optimizations selected hardware platform are
4. Evaluation
4.1. Dataset, Training and Evaluation Setup
Dataset. Our experiments were conducted using real-world sparse matrices sourced from the SuiteSparse Matrix Collec- tion (Davis & Hu, 2011). 1). This dataset has been widely used in previous work (Pal et al., 2018; Hong et al., 2019; Jiang et al., 2020; Gerogiannis al., 2023; Won et al., 2023) and covers broad spectrum of domains, ensuring realistic and comprehensive evaluation of COGNATE's performance. To collect the training dataset, we performed the sparse op- erations SpMM and SDDMM on three distinct hardware platforms: an Intel Xeon Gold 6348 CPU with ITB of RAM, an NVIDIA A100 GPU paired with an Intel Xeon Platinum 8358, and SPADE, simulated sparse accelera- tor with 32 processing elements operating 0.8GHz. To ensure practical feasibility across hardware platforms, the program configuration search space was constrained to 256 configurations for SPADE and approximately 300 configura- tions for Sparse (GPU). We gathered data samples using 1,500 matrices for each hardware platform, with up to 1,000 matrices used for model training under various scenarios and the remainder matrix, we randomly sampled 100 program configurations per hardware platform to have diverse and representative training datasets across all experiments.
that are not applicable zeroed out. Training models on such sparse feature vectors often leads to sub-optimal performance (Figure 4).
To address this limitation, we propose indexing the parame- ters of the heterogeneous component of the program config- urations for each platform a, using low-dimensional latent representations. Specifically, we train an autoencoder AE to learn latent representation 2y for each configuration Cy E Ca.- This is accomplished by determining the value ranges for all parameters of the heterogeneous component in the program configurations, followed by training an autoen- coder to learn an unsupervised embedding of this parame- terization. Once trained, we use the encoder LE in AE which takes each configuration (cz) as input and transforms it into its corresponding latent representation zj. where zj is a fixed-size vector. By compressing configurations from different hardware platforms-each with varying param- eters and ranges-into fixed-size vectors, we standardize the input for hardware-specific optimizations into the cost model. This compression significantly simplifies the model compared feature augmentation, as the cost model now processes fewer input features, reducing its computational
set aside for validation. For each
6

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
Program Configuration Search Space for SPADE. The program configuration search space considered for the
est runtime. We then compared our results to the normalized runtime of the baseline executions, WacoNet with feature augmentation, and WacoNes with feature mapping by cal- culating the geometric mean (geomean) speedups across matrices to quantify COGNATE's overall effectiveness.
SPADE accelerator was derived from combination of key tunable parameters including tiling, synchronization barriers, cache bypassing, and matrix reordering. As sum- marized in Table I, the parameters for barrier insertion, cache bypassing, and matrix reordering are binary (i.e.. ei- ther enabled or disabled). Tiling is controlled by three numerical parameters: the number of row panels, column panels, and the split factor. These values were chosen to resemble those explored in the original SPADE work (Gero- giannis et al., 2023), as those values were expected to show more significant performance deviations for differ- ent sparse matrices. Specifically, we used values for row panels {4, 32, 256, 2048). 4 values for column pan- els {1024, 16384, 65536, NUMMA MUM.MATRIX.CO TRLS (where NUM.MATRIX.COLS depends on the input matrix) and 2 values for the split factor 256).
Pre-training and Fine-tuning Procedure. The matrices for pre-training were randomly selected from the training set while ensuring a balanced representation of their dimen- sions and sparsity. To achieve this, first grouped the matrices into five bins based on the number of rows: S8192, <32,768, 265,536, $131,072, and > 131.072. From each group, we randomly sampled matrices, ensuring the selected subset collectively spanned a diverse range of sparsity lev- els. We empirically demonstrate in Section 4.4 (Figure 11) that training the source model with 100 matrices strikes a good balance. We use this setting for our headline result (Figure 4). Once the source model was pre-trained, we performed few-shot fine-tuning on accelerators using data samples from only 5 matrices, This choice was guided by empirical observations, aiming to strike a good balance be- tween transfer learning effectiveness and the cost of data collection. As shown in Section 4.4 (Figure 12), this set- ting offers the best trade-off between our objectives for accuracy and data collection (detailed in Appendix A). Fur- ther, the same set of matrices were used for evaluating the non-transfer learning baselines, enabling consistent and fair comparisons across all experimental settings.
Although COGNATE is designed to perform under limited data availability, we conducted extensive data collection to rigorously evaluate and justify its effectiveness. This in- cluded range of experiments and ablation studies, some of which required performance data samples from up to 1,000 matrices for training. Altogether, this effort demanded ap- proximately 4 million CPU hours. Despite the constrained nature of the search space (256 program configurations), it took nearly three months to complete the dataset cura- tion, even though the simulations were parallelized across multiple machines. Hence, exhaustively evaluating a larger, unconstrained program configuration space would be com- putationally infeasible. This underscores the need for data- efficient methods like COGNATE, which are designed to operate effectively even under limited data availability.
4.2. Transferability to SPADE
Speedup COGNATE (Top-5): 1.40
Optmal: 1.55
Baseline: 1.0
Geomean Speedup
4
3
Baselines and Implementation. We executed SpMM and SDDMM on CPU, GPU, and SPADE using the respective programming systems introduced in Section 2.2. We used the default optimizations of these programming systems as our baselines. We implemented COGNATE in PyTorch, utilizing MinkowskiEngine (Choy et al., 2019) to handle sparse convolution. Separate models were developed for SpMM and SDDMM to conduct precise performance pre- dictions. We focused on these two sparse operations because they are the only operations currently supported natively by both the SPADE accelerator (Gerogiannis al., 2023) and the Sparse TIR framework (Ye et al., 2023).
2
100
200
300
400
500
600
700
Matrix #
Figure 5: COGNATE per-matrix speedups (SpMM).
Figure 4 illustrates the geomean speedups achieved using multiple techniques: zero-shot inference from the source model (zero-shot), a model trained exclusively on the target hardware using the fine-tuning dataset (no transfer), Wa- coNer with feature augmentation (WACO+FA), WacoNer with feature mapping (WACO+FM), and COGNATE's per- formance for both the top-1 and top-5 (k-best) predicted program configurations. Our results show that COGNATE consistently outperformed other techniques across both sparse operations and hardware platforms. Specifically for SPADE, COGNATE (Top-1) achieved an average speedup of 1.40x for SpMM, reaching 90% of the optimal speedup of 1.55x. When expanding COGNATE (Top-5), it deliv- ered an average speedup of 1.47x, achieving 95% of the
Cost Model Evaluation. We evaluated COGNATE's per- formance on 715 real-world matrices from the SuireSparse Matrix Collection, ensuring that none of the evaluation data samples overlapped with the training set. For each matrix, we predicted the runtime cost across all program configura- tions and selected the 1 and top-5. Then we executed the sparse operations with the selected program configurations on the target platform and identified the one with the short-
7

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
optimal speedup. Note that the optimal speedup was deter- mined by exhaustively evaluating all program configurations within the defined constrained search space for each matrix in the test set, and selecting the fastest configuration per matrix to compute the geometric mean. Similarly, for SD- DMM in SPADE, COGNATE (Top-1) achieved an average speedup of 1.27x and COGNATE (Top-5) achieved an aver- age speedup of 1.39x. This emphasizes COGNATE's ability to consistently find near-optimal program configurations with minimal fine-tuning across multiple sparse operations. The speedup gained for zero-shot inference from the source model was significantly lower than the baseline. In contrast, tinettnniggoo on few data samples from SPADE led to signifi- cant performance gains showing COGNATE's effectiveness in knowledge transfer.
Comparison with Other Transfer Learning Techniques. For comparisons, we modified WacoNer to support feature augmentation and feature mapping, as it is not inherently optimized for heterogeneous transfer. Despite these modi- fications, COGNATE consistently outperformed both. For SpMM on SPADE, WACO+FA had an average speedup of 1.04x, while WACO+FM resulted in slightly higher average speedup of 1.09x. In comparison, COGNATE deliv- ered average speedup of 1. 1.40x, outperforming its closest alternative (WACO+FM) by 28.44%. The sub-optimal per- formance of WACO+FA and WACO+FM can be attributed to the increased sparsity in the feature space due to fea- ture augmentation and their limited capacity to effectively mitigate the heterogeneity.
w/o LE
COGNATE Optimal
LSTM GRU
COGNATE Optimal
w/o FM
Train (Pss( Validation (PSL KKndallls TAU
w/o IFE
Ordered Pair Accracy(OPA)
Geomean Speedup
1.6
Geomean Speedup
1.6
1.55
1.0
Accuracylloss
0.8
0.9
1.4
1.26
1.4
1.32
34
1.40
0.7
0.6
1.2
1.16
1.2
0.5
0.4
1.0
1.0
Epachs
Figure 6: Loss and accuracy during training.
Figure 7: Ablation study of the model components.
Figure 8: Design choices of the predictor.
4.3. Transferability to GPU
4.4. Additional Experiments for SpMM on SPADE
COGNATE is generalizable and is not only applicable to one target hardware platform. To showcase COGNATE's capability, we extended our evaluation to a GPU acceler- ator (NVIDIA A100) (Figure 4). The speedup trends on GPU aligned with those observed on SPADE, reinforcing the effectiveness of COGNATE. COGNATE (Top-1) de- livered an average speedup of 1.03x and COGNATE (Top- 5) yielded an average speedup of 1.17x for SpMM, with the optimal achievable speedup being 1.25x. In compari- son, cUSPARSE SpMM (Naumov et al., 2010) achieved a lower average speedup of 1.01x. For SDDMM, COGNATE (Top-1) resulted in an average speedup of 1.07x, while COGNATE (Top-1) yielded 1.1 speedup, with the op- timal being 1.22x. Zero-shot inference on the GPU was significantly worse compared to Zero-shot for SPADE, with speedups falling well below the baseline. This discrepancy is likely due to the inherent architectural differences be- tween the CPU and GPU architectures. Further, to assess COGNATE's s scalability, we conducted preliminary exper- iments on an end-to-end GNN workload on GPU. Using the 'transient" sparse matrix from our test set (178,866 rows/columns (nodes), 961,368 non-zeros) and GraphS AGE model contigured with three hidden layers and 256 hidden features, COGNATE achieved a 1.30x speedup for inference and 1. .288 speedup for training over the default SparseTTIR (Ye et al., 2023) implementation. These results demon- strate the potential of COGNATE to scale effectively on worldworld workloads and deliver consistent performance.
Speedup Performance. Figure 5 shows the speedups achieved by COGNATE (Top-1) across all evaluated matri- ces. Matrices where the baseline outperformed COGNATE are indicated below the = dotted line. While the baseline outperformed COGNATE on a few matrices, the overall results demonstrate that COGNATE delivered substantial speedups (as high as 5.46x) for the majority of the dataset.
Cost Model Accuracy. Figure 6 shows the accuracy of COGNATE's cost model across training epochs using Pair- wise Ranking Loss (PRL), Ordered Pair Accuracy (OPA), and Kendall's Tau (K-Tau). The steady decline in PRL for both training and validation loss indicates that the model effectively learns to rank program configurations without over-fitting. OPA and K-Tau steadily improved to 0.80 and 0.61, indicating effective training.
Component-Level Contributions. The effectiveness of our cost model relies the inclusion of all components, each contributing uniquely to the overall performance. As illus- trated in Figure 7, the exclusion of individual components leads to noticeable decline in speedups. For example, ex- cluding the input featurizer (IFE) causes a decline from 1.40x to 1.26x. Similarly, omitting the configuration map- per (F leads a further decline to 1.16x, and excluding latent encoder (LE) lowers speedup to 1.01x. This em- phasizes that each component contributes uniquely to the model's high performance, and all need to act synergistically to maximize the benefits of knowledge transfer.
8

---

Figure 12: Impact of number of samples for fine-tuning. COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
PCA
FA
VAE
SCOGNATE
Optimal
NT
1000
COGNATE Optimal
from 1000 matrices) does not necessarily lead to better out- comes. As the size of the training dataset increases, the model becomes overly specialized to the source platform, diminishing its adaptability during fine-tuning. To inves- tigate this effect, we trained source models on datasets of varying sizes (5, 20, 100, 500, and 1,000 matrices) and eval- uated their transferability to our target platform (SPADE) using few-shot tine-tuning on just matrices (Figure 11), Our results show that training data samples from 100 matrices and fine-tuning on SPADE (target) with data samples from 5 matrices produces the best results. In contrast, training the source model with data from 1,000 matrices yields sub-optimal performance due to overfitting to source-specific characteristics. This highlights the importance of carefully selecting the size of the source training dataset to avoid over-ppecialization and minimize the impact of negative transfer.
Geomean Speedup
Geomean Speedup
1.6
155
1.43
14-
1.4
1.40
1.29
1.2
1.2
1.0
1.0
1.0
Figure 9: Selection of auto-
Figure 10: Data overhead w/o transfer learning.
the CPU (source) with
encoders for COGNATE,
Selection of MLP Predictor. As shown in Figure 3, the MLP predictor from WACO's base cost model was retained in our enhanced design. Figure 8 provides comparative analysis of alternative predictors, including LSTM, GRU, and Transformer (TF). The results demonstrate that our proposed cost model design outperforms the alternatives, with the TF predictor achieving the next best performance with 1. 1.36x speedup. These findings highlight that an MLP predictor sufficient to deliver robust performance with limited data. In contrast, the suboptimal performance of the TF predictor can be attributed to the limited dataset, as the high simulation costs associated with emerging hardware make it challenging to collect larger datasets for fine-tuning,
Number of Samples in Fine-Tuning In Figure 12, we show COGNATE's performance as fine-tuning data samples increase. Despite fine-tuning on data from 1,000 matrices, the maximum speedup saturates at 1.42x. We can achieve comparable speedup of 1.40x with data from 5 matrices, which shows the diminishing returns associated with larger datasets. Further, the non- transfer learning setup achieved a marginally higher speedup of 1.43x when using data from 1,000 matrices. However, considering the significant data collection overhead, these marginal improvements are not practically justifiable. To further assess sensitivity to the size of the fine-tuning dataset, we conducted additional ex- periments using and matrices. The resulting speedups were 1.30x and 1.41x, respectively, compared to 1.40x with 5 matrices. While using only 3 matrices led to a noticeable drop in performance. increasing to 7 did not yield a sig- nificant gain but required substantially more data samples, incurring several days of additional data collection time. These results suggest that using 5 matrices strikes a practi- cal balance between data collection cost and performance. while demonstrating that COGNATE is relatively robust to small variations in dataset size.
Selection of Auto-Encoders. Figure 9 shows our inves- tigation into various methods for handling the heteroge- neous components of program configurations. We evaluated choices ranging from conventional feature augmentation (FA) to principal component analysis (PCA), auto-encoders, and variational auto-encoders (VAE). Our findings reveal that auto-encoders were the most effective for capturing het- erogeneous optimizations in a latent space. This was evident from the lower validation loss observed during the training of the auto-encoders to learn the latent representations.
Data Collection Overhead w/o Transfer Learning. Figure 10 shows that without transfer learning, the overhead of data collection becomes significant emerging hardware due to the high costs of running simulations. For example, models trained exclusively on SPADE would require 20x- 200x more target data samples (collected using 100-1000 matrices) to match or surpass the speedups achieved through COGNATE via transfer learning.
5. Conclusion
CPU COGGATEE
CPU 500
Baseline
1000
In this paper, we introduced COGNATE, a novel frame- work to develop data-frugal learned cost models to optimize sparse tensor programs for emerging hardware platforms. COGNATE leverages a unique technique that capitalizes on the homogeneity of input features across different platforms while effectively mitigating heterogeneity. This enables COGNATE to train cost models using low-cost data sam- ples from widely accessible general-purpose hardware (such as CPUs) and then fine-tune them for emerging hardware platforms with few-shot learning. Our results demonstrate that COGNATE is able to achieve near-optimal accuracy while maintaining significant sample efficiency.
CPU Optimal
1000
100
Optimal
Geomean Speedup
1.6
1.55
Geomean Speedup
1.6
1.55
1.40.411 411.421-42
1.40
26
121
1.2
1.2
1.07
1.0
1.0
Figure 11: Impact of nega- tive transfer for fine-tuning.
Impact of Negative Transfer. Figure 11 shows that using a large dataset to train the source model (e.g., data samples
9

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
Acknowledgements
automated (End-to-End) optimizing compiler for deep learning. In 13th USENIX Symposiun on Operating Sys- tems Design and Implementation (OSDI 18), pp. 578-594, 2018a.
The work partially supported by the National Science Foundation Graduate Research Fellowship. ACE, one of the seven centers in JUMP 2.0, a Semiconductor Research Cor- poration (SRC) program sponsored by DARPA, by NSF der the grants CCF 233-3938 and CCF-2316233, by DARPA under the grant D24AP00295-00 and by generous gifts from Qualcomm. Any opinion, findings, and conclusions or rec- ommendations expressed in this material are those of the authors(s) and do not necessarily reflect the views of the NSF or DARPA.
Chen, T., Zheng, L., Yan, E., Jiang, Z., Moreau, T., Ceze, L., Guestrin, C., and Krishnamurthy, A. Learning to opti- mize tensor programs. Advances in Neural Information Processing Systems, 31, 2018b.
Child, R.,
Gray, S., Radford, A., and Sutskever, Gen- erating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
Impact Statement
Choy, C.,, Gwak, J., and Savarese, 4d spatio-temporal convnets: Minkowski convolutional neural networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 3075-3084, 2019.
The goal of this work is to accelerate the execution of sparse tensor programs in the domain of emerging sparse accel- crators through the application of machine learning-based techniques. Experiments demonstrate that our approach exhibits potential in benefiting early-stage accelerator devel- opment by enabling data-efficient design space exploration. There may be potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Dao, T., Chen, B., Liang, K., Yang. J., Song, Z., Rudra, A., and Re, C. Pixelated butterfly: Simple and efficient sparse training for neural network models. arXiv preprint 12.0022 2021.
Daumé III, H. Frustratingly easy domain adaptation. arXiv preprint arXiv:0907.18155 2009.
References
Davis, T. A. and Hu, Y. The university of florida sparse matrix collection. ACM Transactions on Mathematical Software (TOMS), 38(1):1225, 2011.
Aanantharrishnan, S., Abedin, S., Cavé, V., Checconi, F., Du Bois, K., Eyerman, S., Fryman, J. B., Heirman, W., Howard, J., Hur, I., et al. The intel programmable and integrated unified memory architecture graph analytics processor. IEEE Micro, 43(5):78-87) 2023.
Day, O. and Khoshgoftaar, T. M. A survey on heterogeneous transfer learning. Journal of Big Data, 4:1-42, 2017.
Adams, A., Ma, K., Anderson, L., Baghdadi, R., Li, T.- M., Gharbi, M., Steiner, B., Johnson, S., Fatahalian, K., Durand, F., et al. Learning to optimize halide with tree search and random programs. ACM Transactions on Gruphics (TOG), 38(4):1-12, 2019.
Duan, L., Xu, D., and Tsang, Learning with augmented features for heterogeneous domain adaptation. arXiv preprint arXiv: 12XXxv:2000.444 120644660 2012.
Gerogiannis, G., Yesil, S., Lenadora, D., Cao, D., Mendis, C., and Torrellas, J. Spade: A flexible and scalable ac- celerator for spmm and sddmm. In Proceedings of the 50th Annual International Symposium Computer Ar- chitecture, ISCA '23, New York, NY, USA, 2023. Associ- ation for Computing Machinery. ISBN 9798400700958. doi: 10. 10.1144/3559331133388888 org/100 1 /ttp//..1115//5779999 358777111559955 9054.
Baghdadi, R., Merouani, M., Leghettas, M.-H., Abdous, K., Arbaoui, T., Benatchba, K., et al. A deep learning based cost model for automatic code optimization. Proceedings of Machine Learning and Systems, 3:181-193, 2021.
Beltagy, I. Peters, M. E., and Cohan, A. Long- former: The long-document transformer. arXiv preprint arXiv:2004.00 /500 2020.
URL ht htt.s::p//...........t. //doi. /ooi.
Gerogiannis, G., Aananthakrishnan, S., Torrellas, J., and Hur, I. Hottiles: Accelerating spmm with heterogencous accelerator architectures. In 2024 IEEE International Symposium on Biigh-Perffrmacc Computer Architecture (HPCA), pp. 1012-1028. IEEE, 2024.
Bozinovski, S. Reminder of the first paper on transfer learn- ing in neural networks, 1976. Informatica, 44(3), 2020.
Chen, T. and Guestrin, C. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd acm sigkdd inter- national conference on knowledge discovery and data mining, pp. 785-794, 2016.
Graham, B. and Van der Maaten, L. Submanifold sparse con- volutional networks. arXiv preprint arXiv: 170..01307 2017.
Chen, T., Moreau, T., Jiang. Z., Zheng, L.,, Yan, E., Shen, H., Cowan, M., Wang, L.,, Hu, Y., Ceze, L.,, et al. (TVM]: An
10

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
Hegde, K., Asghari-Moghaddaan H., Pellauer, M., Crago, N., Jaleel, A., Solomonik, E.,, Emer, J., and Fletcher, C. W. Extensor: An accelerator for sparse tensor algebra. In Proceedings of the 52nd Annual IEEE/ACM International Symposium on Microarchitecture, pp. 319-333, 2019.
Musio-Martínez, F., Garg, R., Pellauer, M., Abellán, J. L., Acacio, M. E., and Krishna, T. Flexagon: A multi- dataflow sparse-sparse matrix multiplication accelerator for efficient dnn processing. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Vol- ume 3. pp. 252-265, 2023.
Hong, C.,, Sukumaran- Rajam, A., Nisa, I., Singh, K., and Sadayappan, P. Adaptive sparse tiling for sparse matrix multiplication. In Proceedings of the 24th Symposiun on Principles and Practice of Parallel Programming, pp. 300-314, 2019.
Naumov, M., Chien, L., Vandermersch, P., and Kapasi, U. Cusparse library. In GPU Technology Conference, vol- ume 12, 2010.
Jiang, P., Hong, C., and Agrawal, G. A novel data trans- formation and execution strategy for accelerating sparse matrix multiplication on gpus. In Proceedings of the 25th
Neyshabur, B.. Sedghi, H., and Zhang, C. What is being transferred in transfer learning? Advances in neural information processing systems, 33:512-523, 2020.
ACM SIGPLAN symposium on principles and practice of
Pal, S., Beaumont, J., Park, D.-H., Amarnath, A., Feng, S., Chakrabarti, C., Kim, H.-S., Blaauw, D., Mudge, T., and Dreslinski, R. Outerspace: An outer product based sparse matrix multiplication accelerator. In 2018 IEEE International Symposium on High Performance Computer Architecture (HPCA). pp. 724-736, IEEE, 2018.
paraltel programming, pp. 376-388, 2020.
Jin, H., Yue, 7.,, Zhao, Z., Du, Y., Deng, C., Srivastava, N., and Zhang, Z. Vesper: A versatile sparse linear algebra accelerator with configurable compute patterns. IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems, 2024.
Ryu, J. and Sung, H. Metatune: Meta-learning based cost model for fast and efficient auto-tuning frameworks. arXiv preprint arXiv:2102.04199 2021.
Kaufman, S., Phothilimthana, P.. Zhou, Y., Mendis, C., Roy, S., Sabne, A., and Burrows, M. A learned performance model for tensor processing units. Proceedings of Ma- chine Learning and Systems, 3:387-400, 2021.
Sasaki, Y., Takahashi, K., Shimomura, Y., and Takizawa, H. A cost model for compilers based on transfer learning. In 2022 IEEE International Parallel and Distributed Pro- cessing Symposium Workshops (IPDPSW), pp. 942-951. IEEE, 2022.
Kjolstad, F., Kamil, S., Chou, S., Lugato, D., and Ama- rasinghe, S. The tensor algebra compiler. Proc. ACM Program. Lang., 1(OOPSLA), oct 2017. doi: 10.1145/3133011. URL https: //doi cecooooooo. org/10. 1145/3133901.
Shen, Z. Liu, Z., Qin, J., Savvides, M., and Cheng, K.-T. Partial is better than all: Revisiting fine-tuning strategy for few-shot learning. In Proceedings of the AAAI confer- ence on artificial intelligence, volume 35, pp. 2021.
Krizhevsky, A., Sutskever, I., and Hinton, G. F. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25, 2012.
Weiss, K., Khoshgoftaar, T. M., and Wang. D. A survey of transfer learning. Journal of Big data, 3:1-40, 2016.
Langley, P. Crafting papers on machine learning. In Langley, P. (ed.), Proceedings of the 17th International Conference on Machine Learning (ICML 2000). pp. 1207-1216, Stan- ford, CA, 2000. Morgan Kaufmann.
Won, J., Mendis, C., Emer, J. S., and Amarasinghe, S. Waco: Learning workload-aware co-optimization of the format and schedule of a sparse tensor program. In Proceedings of the 28th ACM International Confer- ence on Architectural Support for Programming Lan- guages and Operating Systems, Volume 2, ASPLOS 2023, pp. 920-934, New York, NY, USA, 2023. Associa- tion for Computing Machinery. ISBN 9781 45039911666 doi: 10. 10.11455357569933577777777 org / /0. 110.0.1114/3555550 3575742
Li, Z., Li, J., Chen, T., Niu, D., Zheng, H., Xie, Y., and Gao, M. Spada: Accelerating sparse matrix multiplication with adaptive dataflow. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2. pp. 747-761 2023.
URL ht https: //doi.
Liang, H.. Fu, W., and Yi, F. A survey of recent advances in transfer learning. In 2019 IEEE 19th International Con- ference on Communication Technology (ICCT), pp. 1516- 1523, 2019. doi: 10.1109/CCC4468050200999 201.0 8947077..
Yang, H., Liu, Y., Luan, Z., Gan, L.., Yang. G., and Qian, D. Input-aware sparse tensor storage format selection for optimizing mttkrp. Computer, 56(08):477, aug 2023, ISSN 1558-0814, 14. doi: 10.1109/C.220233227700

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
Ye, Y. and Ji, S. Sparse graph attention networks. IEEE Transactions Knowledge and Data Engineering, 35 (1):905-916, 2021.
Zheng, L., Jia, C., Sun, M., Wu, Z., Yu, C. H., Haj-Ali, A., Wang, Y., Yang, J., Zhuo, D., Sen, K., al. Ansor: Gen- erating Pergh-PPrformacce tensor programs for deep learning. In 14th USENIX symposium on operating sys- tems design and implementation (OSDI 20), pp. 863-879, 2020.
Ye, Z.., Lai, R., Shao, J., Chen, T., and Ceze, L. Sparse- tir: Composable abstractions for sparse compilation in deep learning. ASPLOS 2023, pp. 660-678, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9781450399180
Zheng, L., Liu, R.,, Shao, J., Chen, T., Gonzalez, J. E., Sto- ica, I., and Ali, A. H. Tenset: A large-scale program performance dataset for learned tensor compilers. In Thirty-fith Conference on Neural Information Process- ing Systems Datasets and Benchmarks Track (Round 1), 2021. URL https: htt.ttt//opeeeeeeeeeeeee hwtp//pppeeeeeeeeeeeeeeeeeee.eep net/ forum? id=a T fp8kLuve..
doi: 10.1145/ 3582016.3582047. URL https: //dsii//di..// oro/10. 1145/3582016. 35820447
Zhai, Y., Zhang, Y.. Liu, S., Chu, X., Peng, J., Ji, J., and Zhang, Y. Tlp: A deep learning-basee cost model for tensor program tuning. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume pp. 833 845, 2023.
Zhuang. F., Qi, Z., Duan, K., Xi, D., Zhu, Y., Zhu, H., Xiong, H., and He, Q. comprchensive survey on transfer learning. Proceedings of the IEEE, 109(1):43-76, 2020.
12

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
A. Problem Formulation
In this work, our aim is to build accurate learned cost models for emerging hardware platforms to enable faster identification of optimal program configurations. A key challenge the need to maximize the accuracy of the cost model (aecuracy objective) while using as few expensive (i.e. collected through simulation) data samples as possible (data collection objective). We first formalize the program optimization objective and then tie it with the cost model objectives.
A.1. Program Optimization Selection
The goal of program optimization in sparse tensor programs to select the optimal program configuration for given hardware platform and input sparsity pattern from the total configuration space. Let configuration space C. be the set {ci, ......, J J of all valid program configurations for a given hardware platform (m, E Z+). For example, for CPU, a valid configuration from CCpU is a tuple of program configuration parameters for loop strip-mining, loop reordering, and format reordering (Table 1). The optimal program configuration minimizes the execution time of a sparse tensor program. For an input sparse matrix (sparsity pattern) M, the optimal program configuration on platform a can be given as, e" arg ec. (M,M, (M,cs) es), where Ta is the execution time function for platform a (ground truth runtime). The execution time for the optimal program configuration is given by
A.2. Cost Model Performance and Data Efficiency Objectives
We approximate the ground truth runtime Ta using learned cost models. Usually, these cost models are trained with one objective: achieve high accuracy. However, due the high cost of simulation in emerging hardware, we also want to minimize the amount of data samples required from these platforms for model training. We formalize these two objectives as follows.
Data Collection Objective (DCE), Let D, = {M,, c.), i E ma , E Z+} be the dataset collected from hardware platform a, and let Sa represent the average cost of collecting a single data sample from the platform. Our objective is to minp, Bg x |Dal.
Accuracy Objective. Let C.a (which approximates Ta) be the learned cost model trained on dataset Da. If the best program configuration returned by the cost model (cCM,) ) has an actual execution time tCM., our objective is to min itcm 1. where t* is the execution time for the optimal configuration. For set of input sparse matrices (M1, M2..... Mr). our objective can be extended to minimizing the Absolute Percentage Error (APH) across all matrices:
MCM. - tM
APE
X 100
where tCas. denotes the execution times for the predicted best program configuration for the input sparse matrix M, and t5, denote the optimal program configurations for the same matrix.
A.3. Evaluations for Cost Model Objectives
To evaluate the cost model objectives, we conducted the following experiments for SpMM on SPADE. For simplicity in the calculations, we set fCru and ASPADE 1000. However, CPU execution typically takes milliseconds, while SPADE execution can extend up to two weeks. We explored 11 distinct models across four different categories, differentiated by the number of data samples they were trained on, while the cost model architecture remained the same. Category I consists of models (NT d) trained exclusively on data samples from d matrices executed on SPADE. Category II includes transfer-learned models (TL d). which were pre-trained with data samples from 100 matrices on CPU (10,000 data samples) and then fine-tuned on SPADE with data samples from matrices. Category III consists of models (CPU d) pre-trained with varying numbers of data samples from d matrices on CPU and then fine-tuned on data samples from 5 matrices on SPADE. Finally, we did zero-shot inference (Zero-Shot) from model pre-trained on CPU with data samples from 100 matrices without additional fine-tuning on SPADE.
Models trained exclusively on SPADE data samples (NT d) generally exhibit increasing speedup and decreasing APE as the number of SPADE data samples increases. For example, NT 1000, trained on 100,000 SPADE data samples, achieves the highest speedup of 1.43 and an APE 7.06. However, the data collection overhead for these models rises significantly with the number of SPADE samples, making the use of them impractical due to the long simulation times. In contrast, the TL.
13

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
models, which are pre-trained CPU data and fine-tuned on SPADE samples, demonstrate excellent balance between speedup, APE, and DCE. TL.5 model, for instance, delivers competitive speedup of 1.40 and low APE of 7.28, while maintaining an excellent DCE of 0.51.
Model
Data Samples
COGNATE (Top-1) Speedup
APE
DCE/10 110
CPU
SPADE
5
500
1.29
15.02
0.50
NT 100
10000
1.38
9.42
10.00
NT 1000
100000
1.43
7.06
100.00
5 (CPU 100)
10000
500
1.40
9.58
0.51
TL 100
10000
10000
1.41
8.74
10.01
TL 1000
10000
100000
1.42
7.28
100.01
CPU 5
500
500
1.07
27.80
0.50
CPU 20
2000
500
1.21
19.35
0.50
CPU 500
50000
500
1.36
16.34
0.55
CPU 1000
100000
500
1.19
36.00
0.60
Zero-Shot (CPU)
10000
0.71
46.22
0.01
Table 2: Comparison of cost model performance with varying data samples from CPU and SPADE.
A.4. Learning Objective
Our objective is to train cost model that effectively learns to identify program configuration that minimizes the runtime of sparse operation for given sparsity pattern. To achieve this, we begin by training our cost model to learn the relative rankings of program configurations during execution, enabling it to accurately identify optimal configurations based on their performance. This objective improves robustness to noise and runtime scale variance, which are common in early-stage accelerator performance data, as the model focuses on preserving relative orderings. This also enables us to efficiently integrate our cost model with search technique to efficiently select the top-k (k-best) program configurations from the configuration space. Furthermore, prior work (Kaufman et al., 2021) has shown that training with ranking loss significantly improves a model's ability to identify optimal configurations. We use the pairwise ranking loss our learning objective (implemented using margin ranking loss) to rank program configurations based on their true performance differences. For a given input matrix M, the pairwise ranking loss (C) across all program configuration pairs can be defined as L= E(er, (2)) max(0, 1 - (fM,q, - fM.c))) )) Spnes 8ine = sign(t.,ec - fas,es) where fMe,, and Ê, are the predicted scores for configurations с₁ and C2, respectively; (are, and (M represent their actual runtimes; and ôene signifies the true performance difference where sign(a) returns x > 0, -1 if < O, and 0 if = This ensures that the model penalized when the predicted ranking does not align with the true ranking. By minimizing this loss (C). COGNATE improves its ability to accurately rank and identify the top-k program configurations. This also contributes to achieving accuracy objective (Section A.2).
B. Code Optimizations Across Hardware Platforms
Loop strip-mining:
Breaks down large software loops into smaller segments to optimize cache utilization.
Loop reordering: Adjusts the execution order of loops to improve cache efficiency. Typically, it is applied after loop strip-mining.
Format reordering: Reorganizes the data structure layout of sparsity patterns in memory to optimize memory access patterns
Parallelization: Distributes tensor computations across multiple threads or processors run tasks simultaneously.
Loop binding: Assigns specific loop iterations to threads for parallel processing.
Loop unrolling:
Executes multiple iterations of a loop in single iteration, reducing loop control overhead and boosting execution speed.
14

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
Tiling: Decomposes matrix into smaller blocks to optimize data reuse in the local memory and improve cache efficiency.
Applying a barrier would ensure all threads finish processing their current tile (synchronized) before progressing to the next stage.
Barrier:
Cache bypassing: Capability of bypassing caches to to reduce cache pollution.
Matrix reordering: Enhances data locality by reordering the input matrix.
C. Generalizability
Our approximated code mappings are designed to enable well-established code optimizations that remain effective across emerging hardware platforms. For instance, we expect loop transformations (e.g. loop strip-mining, loop reordering, tiling, etc) to be implemented regardless of the underlying hardware platform although the exact implementation may be slightly different. Any hardware-specific optimizations (code optimizations that cannot be mapped) are included in the heterogeneous component using the latent encoder. Since the dimension of the latent embedding is fixed, we can effectively finetune using already pre-trained model. To elaborate, let us have qualitative discussion and explore the intuition behind incorporating approximate mapping of these loop transformations into sparse accelerators, using Intel PIUMA (Gerogiannis et al., 2024; Aananthakrishnan et al., 2023) and Vesper (Jin et al., 2024) as examples. These mappings are conceptually aligned with those we applied to SPADE and GPU, highlighting the general applicability of our approach across diverse hardware backends.
Intel PIUMA (Gerogiannis et al., 2024; Aananthakrishnan et al., 2023) is a configurable accelerator that has a RISC ISA making CPU-programmable. This enables it to employ code optimizations that are available in CPUs. Hence, it is possible to implement SpMM and SDDMM sparse operations (kernels) with code optimizations such loop reordering with one-to-one mapping. Similarly, loop strip-ming and tiling can be mapped. However, similar to SPADE, where we accounted for the "barrier" optimization in the mapping process, one would need to consider the PIUMA "scratchpad reuse".
Vesper (Jin et al., 2024) is another reconfigurable accelerator designed for sparse computations, supporting three dataflow models implemented through distinct loop traversal orders. While the authors refer to the use of "tiling," they do not provide implementation details, source code, or description of the tile size selection mechanism. Based on standard tiling practices, we can approximate Vesper's 's approach using loop stripping and loop reordering within our representation.
Intel PIUMA is proprictary, and Vesper's 's source code was not available. This made it infeasible to test our hypothesis on these accelerators. Hence, our current evaluation focuses only on two examples (SPADE and NVIDIA A A100) primarily due to practical constraints. However, it should be emphasized that COGNATE was designed with hardware-agnostic principles in mind. We believe that as a wider range of accelcrators becomes accessible to the research community, and as sparse compilation frameworks like SparseTIR (Ye et al., 2023) evolve, COGNATE can be extended with minimal changes.
Assuming these accelerators were available, the data collection process would still be highly time-consuming, likely requiring millions of machine hours gather sufficient data for training. validation, and testing. For example, collecting performance data or training, validation, and testing across all matrices and experimental settings for SPADE required approximately 4 million CPU hours. Despite parallelizing experiments across multiple machines, each with 64 CPU cores, this process spanned nearly three months. While extending the evaluation to additional hardware platforms remains an important direction, it was beyond the practical scope of this work given resource constraints.
Further, frequent changes in emerging hardware may require updates configuration mappings and fine-tuning, this challenge is significantly mitigated by our approach in COGNATE. As long as the entire kernel does not change or the newly introduced optimizations are heterogeneous, updating the mappings is relatively straight forward. However, relying solely on simulations would require rerunning them for large number of configurations each time a change is made, resulting in significant computational and time costs. In contrast, our transfer learning- approach significantly reduces the cost and time of running simulations. By collecting only a few data samples and finettunig the model, we can efficiently adapt to hardware changes without the need for extensive simulations. Hence, this approach not only reduces maintenance complexity but also accelerates the design process, making it more feasible to handle frequent and timely updates in emerging hardware.
15

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
D. Cost Models for Early-Stage Sparse Accelerator Design
With the flexibility of recent accelerators to have software programmable kernels (Gerogiannis et al., 2023; Jin et al., 2024; Gerogianmis et al., 2024), integration of cost models and heuristies into the DSE pipeline has become an up and coming area (Gerogiannis al., 2024; Jin al., 2024). For example, Vesper is a recent work that had integrated an analytical model to configurable sparse accelerator to enable higher throughput (Jin et al., 2024). HotTiles is another work that uses an analytical model to predict the performance of different accelerator processing elements (PEs) that accommodate intra-matrix heterogeneity (Gerogiannis et al., 2024). Further, in HotTiles, the authors acknowledged that a more accurate model could have enabled making better design decisions during the early stages. We believe that our proposed data-driven cost model framework, COGNATE, addresses this gap (resulting in speedups close to optimal) while complementing expert-driven strategies to enable more informed and better design decisions. This would effectively replace the analytical approaches with data driven approach. The primary overhead associated with our approach arises from the need to gather data points to fine-tune the cost model. This overhead is minimal compared to the effort required for an expert to iteratively optimize kernel for sparse workloads, where kernel performance is highly input-sensitive due to diverse sparsity patterns.
E. An Example of Code Optimization Notations Used in Approximate Mappings
Here, we provide an additional explanation and an illustrative example to clarify the notation used in the code op- timization mapping functions presented in Section 3.2. These are designed to approximate how high-level schedule configurations in SPADE are translated into low-level loop representations in CPU. The following example demon- strates how sparse matrix-matri multiplication (SpMM) configuration in SPADE is mapped into its corresponding loop-level representation using the defined notation. Consider the following high-level configuration for the SpMM oper- ation in SPADE: name, row.panels, col usn, split, barrier, bypass, reorder, time 1 144, 4, 1024, 1, 0, 0, 0, 38. 83392. Here, row panels, column panels, and sp it define the tiling strategy, while the binary flags barrier, bypass, and reorder indicate the use of additional code op- timizations. Using our mapping functions, this configuration tation: name, i-split, j-split, k-split, loop_l,
mapped into the following loop-level represen- 1oop-7, barrier, bypass, reorder, time- = 144, 4, 1024, 32, 6, 7, 2, 4, 1, 3, 5, 0, 0, 0, 38. 83592. In this mapped form, the tiling parameters are converted to isspit, j-split, and k_split, which define how the loop indices are partitioned across the three dimensions. The sequence loop 1 through loop 7 encodes the execution order of the nested loops, and the binary flags are retained to preserve platform-specffic scheduling decisions. This example demonstrates how our framework captures key aspects of tiling structure, loop ordering, and other scheduling optimizations.
F. Hyperparamters
Table 3: Hyperparameters for model training/fin--tuiing
Hyperparameter
Value
Learning Rate
0.0001
Batch Size
32
Optimizer
Adam
Number of Epochs
100
Loss Function
MarginRankingL.oss
Table 4: Hyperparameters for the autoencoders
Hyperparameter
Value
Learning Rate
0.001
Batch Size
32
Optimizer
Adam
Number of Epochs
1000
Loss Function
MSE
16

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
Table 5: Composition of layers in the Input Featurizer (IFE)
Layer
Description
Layer
MinkowskiConvvllttin (in channels, 32, kernel size=5)
Layer
MinkowskiConnvlution (32, 32, kemnelsize==3
Layer
MinkowskiConnvolttion (32, 64, kernelsiiz==3) MinkowskiMMaxPoolin
Layer
MinkowskiConnvllttin (64, 64, kemnelsize--)
Layer
MinkowskiiConvllttinn (64, 64, kernel_size=3)
Layer
Minkow>iiConvvolttionn (64, 128, kernelsiize=3) Mino
Layer
MinkowskiCConvolution (128, 128, kernel
Layer
MinkowskiiConnvllttion (128, 128, kernel_size=3)
Layer
Minkowsk(Convolution2222, 256, kernel Lize=s)Miiikwwwsiiaaooooing
Layer 10
MinkowskiCoonvolltion (256, 256, kernel size=3)
Layer 11
MinkowskiiConvolltiion (256, 256, kernel_size=3)
Layer 12
MinkowskiConnolutionn (256, 256, kernel kize=)_size=3)
Global Pooling Layer
MinkowskiiGiohallAyyPoolin
Table 6: Composition of layers in the Predictor (P)
Component/Layer
Input Size
Output Size
Matrix Embedding (x)
128
128
Configuration Embedding (y)
53
64
Latent Embedding (z)
64
64
Concatenation (xyz)
128
64
192
Predictor Layer 1
192
128
Predictor Layer 2
128
64
Predictor Layer 3
64
17

---

COGNATE: Acceleration of Sparse Tensor Programs on Emerging Hardware using Transfer Learning
G. Additional Results
Geomean Speedup
Speedup
COGNATE (Top-5); 1.47 Baseline: 1.0
5
Optimal: 1.55
4
100
200
300
400
500
600
700
Matrix
Figure 13: COGNATE (Top-5) per-matrix speedups (SpMM)
Speedup
COGNATE TTop1):: 1.27
Geomean Speedup
4.0
Optimal: 1.44--- Basellie: 1.0
3.5
3.0
2.5
2.0
1.5
1.0
0.5
400
500
600
700
Matrix
Figure 14: COGNATE (Top-1) per-matrix speedups (SDDMM)
Geomean Speedup
Speedup
COGNATE (Top-5): 1.39
4.0
Optimal: 1.44
Baseline: 1.0
3.5
3.0
2.5
2.0
0.5
100
400
500
600
700
Matrix
Figure 15: COGNATE (Top-5) per-matrix speedups (SDDMM)
18


Reviews:
[{"review_id": "hwCLyRIju6", "reviewer": "Reviewer_cfbi", "summary": "The paper introduces COGNATE, a novel framework for developing learned cost models to optimize sparse tensor programs on emerging hardware platforms. COGNATE leverages transfer learning to adapt cost models from general-purpose hardware (e.g., CPUs) to specialized accelerators with minimal fine-tuning data.\n\nMain Contributions:\n- Transfer Learning Approach: COGNATE uses a pre-trained model on CPU data and fine-tunes it on sparse accelerators, significantly reducing the need for expensive simulation data.\n- Homogeneity and Heterogeneity Handling: It maps similar code optimizations across platforms and uses autoencoders to encode heterogeneous components, allowing for efficient knowledge transfer.\n- Performance: COGNATE achieves average speedups of 1.47× for Sparse Matrix-Matrix Multiplication (SpMM) and 1.39× for Sampled Dense-Dense Matrix Multiplication (SDDMM) on the SPADE accelerator.\n\nUpdate after rebuttal:\n\nI would like to sincerely thank the authors for taking the time to respond to all the issues I raised. However, as I pointed out in the initial review, the design and evaluation are quite limited with SPADE and NVIDIA A100 GPU. The authors claimed that just focusing on two examples (SPADE and NVIDIA A100) is primarily due to practical constraints, such as lack of access to closed-source platforms. However, there are plenty of better platforms than closed-source platforms that can be accessed, just to name a few: NVIDIA Jetson AGX Orin, NVIDIA H20 & H100 GPU. Moreover, from the authors' replies, the effectiveness of the proposed COGNATE method is highly related to the hyper-parameter settings. For example, the size of the dataset will significantly influence the transfer effect. So, it brings the concern of the generalization of the COGNATE method. Finally, the Sparse Tensor Core features may change significantly for different accelerator hardware platforms. For example, H100 GPU further supports the new FP8 data type, and the combination of Tensor Memory Accelerator (TMA). The COGNATE method also fails to address how to fit the 2:4 structured sparse pattern, which is the most important feature introduced in the NVIDIA A100 sparse Tensor Core. With all these concerns, I still hold my initial score. So, the current version may not meet the acceptance threshold for ICML.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for recognizing the significance of the problem we address and for acknowledging the contributions of our work. We are especially grateful for the time and effort you invested in providing a detailed and thoughtful summary of the paper’s strengths and areas for improvement. We hope that our following response addresses suggestions and questions.\n\n**Q1 (Question) & W1 (Weakness)**: We thank the reviewer for this thoughtful question. In Appendix Section C (Generalizability), we provide an initial qualitative discussion and intuition on how COGNATE can be generalized. Our current evaluation focuses on two examples (SPADE and NVIDIA A100) primarily due to practical constraints, such as lack of access to closed-source platforms (PIUMA and Vesper). Even if we had access, the data collection for these accelerators would be highly time-consuming, requiring millions of machine hours. For instance, collecting performance data for SPADE required approximately 4 million CPU hours. Hence, while extending our evaluation to more hardware platforms would be desirable, it was not feasible given the computational resources and simulators we had access to. We believe that as a wider range of accelerators becomes accessible to the research community, COGNATE can be extended with minimal changes. We will further clarify this point in the final version of the paper. \n\n**W2 & W3**: We thank the reviewer for pointing this out. In the final version, we will improve the content to enhance clarity and include more details about hyperparameter tuning.\n\n**Q2**: We thank the reviewer for this insightful question. We explored the effects of negative transfer by training the source model using data samples with 5, 20, 500, and 1000 matrices (Figure 11). While a larger dataset improves generalization in the source domain, it may also cause over-specialization to the source platform. This overfitting hampers fine-tuning, as specialized features may not transfer well to the target platform. Our empirical findings show that training the source model on 100 matrices strikes a good balance, capturing useful patterns while maintaining generality for transfer. This reduces the effort needed during fine-tuning, leading to faster convergence and better performance. This is especially valuable given the high cost of collecting data for SPADE.\n\n**Q3**: We thank the reviewer for this important question. Our current evaluation focuses on SpMM and SDDMM because these are the sparse operations currently natively supported by the SPADE and SparseTIR. That said, SpMM and SDDMM serve as foundational blocks for complex programs. To evaluate COGNATE’s scalability, we conducted preliminary experiments during the rebuttal period on an end-to-end GNN workload running on a GPU, using the ‘transient’ sparse matrix from our testset as input (178,866 nodes with 961,368 non-zeros) and GraphSAGE model. The model was configured with 3 hidden layers, 256 hidden features. COGNATE achieved notable speedups over the default SparseTIR implementation, with a 1.30× speedup for inference and 1.28× for training, demonstrating the scalability of COGNATE. We will include this result in the final version of the paper and highlight it as a key direction moving forward.\n\n**Q4**: We appreciate the reviewer raising this point. While changes in emerging accelerators may necessitate updates to configuration mappings or model parameters, COGNATE significantly reduces this burden by relying on lightweight fine-tuning, rather than requiring retraining from scratch. As long as the overall structure of the sparse tensor program remains consistent, COGNATE can quickly adapt by using a small number of new performance samples. In contrast, traditional cost model construction approaches would require re-evaluating a large number of configurations. We provided further elaboration on this in under Generalizability (Section C) in the Appendix.\n\n**Q5**: We appreciate the reviewer’s interest in the training objective. Our goal during model training is not to predict the absolute runtime of a sparse tensor program configuration, but rather to rank candidate configurations by their relative performance, enabling the selection of the best-performing ones. To this end, we adopted a pairwise ranking loss, which is more aligned with the optimization task than pointwise losses. This objective improves robustness to noise and runtime scale variance, which are common in early-stage accelerator performance data, as the model focuses on preserving relative orderings. Furthermore, prior work (Kaufman et al., 2021) has shown that training with ranking loss significantly improves a model’s ability to identify optimal configurations. We will clarify this motivation and include these details in the final version to make the learning objective and its impact more explicit. Further, if space permits, we will move the loss equation (in Appendix Section A.4) to the main text of the paper.", "claims_and_evidence": "The claims made in the submission are generally supported by clear and convincing evidence, but there are a few aspects that could be scrutinized further:\n1. Performance Comparisons:\n- Evidence: The paper provides extensive comparisons with existing techniques like WACO+FA and WACO+FM, showing that COGNATE outperforms them by a significant margin (28.44% improvement for SpMM on SPADE)\n- Potential Issue: While the comparisons are thorough, it would be beneficial to see more detailed analysis on why COGNATE performs better, especially in terms of its ability to handle heterogeneity and its data efficiency\n\n2. Data Efficiency:\n- Evidence: COGNATE is shown to achieve comparable performance with only 5% of the data required by accelerator-specific models, which is a significant reduction in data collection overhead\n- Potential Issue: The paper could further elaborate on how this efficiency is achieved, particularly in terms of the autoencoder's role in compressing heterogeneous features and the impact of using a reduced number of layers in the cost model\n\n3. Generalizability:\n- Evidence: COGNATE demonstrates its applicability across different hardware platforms, including SPADE and NVIDIA A100 GPU, with notable speedups in both cases\n- Potential Issue: While the results are promising, additional evaluations on more diverse hardware platforms would strengthen the claim of generalizability", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria in the paper make sense for the problem of optimizing sparse tensor programs on emerging hardware platforms. The paper addresses a critical challenge in optimizing sparse tensor programs, which are essential in deep learning and graph analytics. The use of real-world sparse matrices from the SuiteSparse Matrix Collection provides a comprehensive and realistic evaluation setup.", "theoretical_claims": "The paper on COGNATE does not present formal proofs for its theoretical claims. Instead, it focuses on empirical evaluations and algorithmic design to support its contributions. However, there are no explicit mathematical proofs provided in the paper for theoretical claims like:\n- Optimality of the Transfer Learning Approach: The paper demonstrates empirically that COGNATE outperforms other transfer learning techniques but does not provide a theoretical proof of optimality.\n- Effectiveness of Latent Encoding: The use of autoencoders to handle heterogeneous components is shown to be effective in practice, but there is no formal proof of its theoretical advantages over other methods like feature augmentation.\n- Data Efficiency: While the paper shows that COGNATE achieves comparable performance with significantly less data, there is no formal proof that this approach is optimal in terms of data efficiency.", "experimental_designs_or_analyses": "The experimental design and analyses in the paper on COGNATE appear to be sound, but there are a few aspects that could be scrutinized further:\n- The evaluations on more diverse hardware platforms (e.g., other specialized accelerators, FPGAs) could strengthen the claim of generalizability.\n- Further analysis on how the number of fine-tuning samples affects performance could provide deeper insights into the limits of COGNATE's data efficiency.\n- Including other metrics (e.g., energy efficiency, memory usage) could offer a more holistic evaluation of COGNATE's benefits.", "relation_to_broader_scientific_literature": "The key contributions of the paper on COGNATE are closely related to the broader scientific literature in several ways:\n1. Transfer Learning for Cost Models:\nPrior Work: Transfer learning has been successfully applied in various domains to reduce data requirements for target tasks (Weiss et al., 2016; Zhuang et al., 2020). In program optimization, transfer learning has been used to adapt cost models across similar hardware platforms (Sasaki et al., 2022; Zheng et al., 2021).\n\n2. Handling Heterogeneity in Transfer Learning:\nPrior Work: Existing heterogeneous transfer learning techniques, such as feature augmentation and feature mapping, have limitations when dealing with diverse program configurations across different hardware platforms (Daumé III, 2009; Duan et al., 2012).\n\n3. Optimization of Sparse Tensor Programs:\nPrior Work: Optimizing sparse tensor programs is crucial for deep learning and graph analytics, with techniques like TACO (Kjolstad et al., 2017) and WACO (Won et al., 2023) providing significant performance improvements.\n\n4. Data Efficiency in Early-Stage Hardware Development:\nPrior Work: The high cost of collecting large datasets for emerging hardware platforms is a significant challenge (Gerogiannis et al., 2023).", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "The paper could provide more discussion on the broader impact of COGNATE beyond the specific hardware platforms evaluated (SPADE and NVIDIA A100 GPU). Demonstrating its applicability to a wider range of accelerators or scenarios could enhance its significance.\n\nSome sections, particularly those detailing the mapping functions and autoencoder training, could benefit from additional illustrations or step-by-step explanations to improve clarity for readers unfamiliar with these techniques. Additionally, the paper assumes a strong background in sparse tensor programs and transfer learning, which might limit accessibility for readers from other domains.\n\nThe paper could provide more details on the hyperparameter tuning process and how specific hyperparameters were chosen. This would help in reproducing the results and understanding the sensitivity of COGNATE to different hyperparameter settings.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. How do you envision extending COGNATE to support a broader range of emerging hardware platforms beyond SPADE and NVIDIA A100 GPU? Are there specific challenges or modifications needed for other types of accelerators?\n\n2. You mention that using a moderate-sized dataset for the source model helps mitigate negative transfer. Can you provide more insights into how the size of the source dataset affects the fine-tuning performance on target platforms?\n\n3. As the complexity of sparse tensor programs increases, how does COGNATE's performance scale? Are there any plans to address more complex operations or larger-scale programs?\n\n4. Given the rapid evolution of hardware, how does COGNATE adapt to changes in hardware architecture or new optimizations introduced in emerging accelerators?\n\n5. Can you elaborate on the pairwise ranking loss used in training the cost model? How does this objective function contribute to the model's ability to identify optimal program configurations?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FWDgNXLBo4", "reviewer": "Reviewer_eqda", "summary": "The submission introduces COGNATE, a novel framework designed to optimize sparse tensor programs on emerging hardware accelerators using machine learning-based cost models. It addresses the challenges of optimizing sparse tensor programs, such as Sparse Matrix-Matrix Multiplication (SpMM) and Sampled Dense-Dense Matrix Multiplication (SDDMM), on early-stage accelerators where performance is sensitive to sparse input variations and data collection via simulators is costly. COGNATE leverages transfer learning by pre-training cost models on inexpensive CPU data and fine-tuning them with minimal accelerator-specific data (5% of typical requirements), achieving high sample efficiency.\n\nKey contributions include: (1) techniques to segregate program configurations into homogeneous (mapped via approximate code optimization mappings) and heterogeneous (encoded via autoencoders) components, enabling effective knowledge transfer across hardware platforms; (2) a data-frugal cost model framework that modifies the WACO architecture for better transferability, incorporating a configuration mapper, enhanced input featurizer, latent encoder, and predictor; and (3) extensive evaluations demonstrating COGNATE’s effectiveness. Main results show COGNATE achieves average speedups of 1.47× (up to 5.46×) for SpMM and 1.39× (up to 4.22×) for SDDMM on the SPADE accelerator, and 1.17× (up to 1.61×) for SpMM and 1.15× (up to 1.49×) for SDDMM on an NVIDIA A100 GPU, outperforming existing transfer learning methods by 28.44% on SPADE. These findings highlight COGNATE’s ability to deliver near-optimal performance with minimal data overhead, enhancing design space exploration for emerging sparse accelerators.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for recognizing the significance of our contributions. We are especially grateful for the time and effort you dedicated to providing such a detailed and thoughtful review. We also appreciate the additional references you shared, which we'll include in the final version. We hope that our following response addresses your suggestions and questions.\n\n> ***Other Comments*** \\\n> ... key hyperparameters for autoencoders … to improve reproducibility.\n\nWe thank the reviewer for pointing this out. The hyperparameters for the autoencoder training are:\n\n- Learning Rate: 0.001\n- Loss: MSE\n- Optimizer: Adam\n- Batch Size: 32\n- Epochs: 1000\n\nWe will include this information in a table similar to Table 3.\n\n> ***Question 01*** \\\n> ... how sensitive are results to slight variations (e.g., 3 or 7 matrices)?\n\nWe thank the reviewer for this thoughtful question. The choice of 5 matrices for fine-tuning was primarily guided by empirical observation and the need to balance transfer learning effectiveness with the cost of data collection. As mentioned in the paper, “the matrices were randomly selected from the training set while ensuring a balanced representation of their dimensions and sparsity.” To do this, we first divided the available matrices into five groups based on input size [8192, 32768, 65536, 131072, and > 131072], then randomly sampled a matrix from each group while ensuring the selected subset captured a reasonable range of sparsity. We used the same randomly selected matrices for the non-transfer learning baselines to ensure consistency and fair comparison across experimental settings.\n\nIn response to the reviewer’s suggestion, we conducted additional experiments for SpMM using 3 and 7 matrices to assess sensitivity to small variations in fine-tuning set size. The observed top-1 speedups were (compared to 1.40×);\n\n- 3 matrices: 1.301×\n- 7 matrices: 1.405×\n\nThese results suggest that COGNATE is relatively robust to small variations in the size of the fine-tuning dataset. While using only 3 matrices results in a slight drop in performance, the model still achieves meaningful gains over the zero-shot baseline. Finally, we acknowledge that no formal statistical power analysis was used in selecting the fine-tuning size. We view this as a promising direction for future work.\n\n> ***Weakness 01*** \\\n> While innovative, it heavily builds on WACO …, potentially limiting its perceived novelty in the ML optimization space.\n\nWe thank the reviewer for this observation. We agree that our work builds on top of WACO, the current state of the art in ML-based sparse autotuning. However, the key contributions of COGNATE are distinct and extend beyond WACO in several important ways. While we adopt WACO’s search space and program representation due to their relevance to sparse tensor programs, our primary focus is cross-platform generalization and efficient model transfer, which are orthogonal to WACO’s contributions. Specifically, COGNATE’s framework leverages the homogeneity of input features across hardware platforms while mitigating heterogeneity to efficiently fine-tune learned cost models for accelerating sparse operations on emerging hardware. As part of future work, we are actively exploring the application of COGNATE to other domains beyond sparse tensor programs and WACO’s framework.\n\n> ***Weakness 02*** \\\n> Generalizability to untested accelerators (e.g., PIUMA, Vesper) is speculative, ...\n\nWe acknowledge the reviewer’s concern. Our current evaluation focuses on two examples (SPADE and NVIDIA A100) primarily due to practical constraints, including lack of access to closed-source platforms like PIUMA and Vesper. Even if we had access, the data collection process for these accelerators would be highly time-consuming, likely requiring millions of machine hours to collect sufficient data for training, validation, and testing. For instance, collecting performance data for training, validation, and testing for SPADE required approximately 4 million CPU hours.  Although we parallelized the required experiments across 32 machines, each with 64 CPU cores, this process took us nearly three months.  Hence, while extending our evaluation to even more hardware platforms would be desirable, it was not feasible given the computational resources and simulators we had access to.\n\nHowever, we emphasize that COGNATE was designed with hardware-agnostic principles in mind. We believe that as a wider range of accelerators becomes accessible to the research community, and as sparse compilation frameworks like SparseTIR evolve, COGNATE can be extended with minimal changes. We will further clarify this point in the final version of the paper. That being said, in the final version, we will tone down our generalization claims and more clearly state that the quantitative demonstration of COGNATE’s even broader applicability (beyond the two hardware accelerators we evaluated) remains an important direction for future work.", "claims_and_evidence": "The claims in the submission are generally supported by clear and convincing evidence, including detailed experimental results, figures, and tables comparing COGNATE’s performance against baselines and alternative methods. The main findings—speedups of 1.47× for SpMM and 1.39× for SDDMM on SPADE, and 1.17× and 1.15× on A100 GPU—are backed by evaluations on 715 real-world matrices from the SuiteSparse collection, with geomean speedups, per-matrix speedups (Figures 5, 13-15), and accuracy metrics (Figure 6) provided. The claim of data efficiency (using 5% of typical data) is substantiated by comparisons of data collection overhead (Figure 10) and fine-tuning sample analysis (Figure 12). Ablation studies (Figure 7) and component design choices (Figures 8-9) further support the algorithmic innovations.\n\nHowever, two claims could be seen as less robustly evidenced:\n\n- Generalizability to Other Accelerators: The paper claims COGNATE is generalizable beyond SPADE and A100, referencing Intel PIUMA and Vesper (Section C). However, no experimental data is provided for these platforms due to proprietary restrictions or unavailable source code, weakening this claim with speculative rather than empirical support.\n\n- Near-Optimal Accuracy: The assertion of \"near-optimal accuracy\" (e.g., 95% of optimal speedup on SPADE) relies on comparing COGNATE’s top-5 predictions to an optimal baseline derived from exhaustive configuration testing. While results are strong, the lack of a statistical significance test or error bounds around the 95% figure undermines the claim.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria in COGNATE are well-suited to the problem of optimizing sparse tensor programs on emerging hardware accelerators. The transfer learning approach—pre-training on CPU data and fine-tuning with minimal accelerator data—addresses the high cost of simulator-based data collection and input sensitivity of sparse operations like SpMM and SDDMM. Key components (configuration mapper, latent encoder, enhanced input featurizer) logically tackle the heterogeneity and data efficiency challenges outlined.\n\nThe evaluation criteria, including geomean speedups, top-1/top-5 predictions, and accuracy metrics (Pairwise Ranking Loss, Ordered Pair Accuracy, Kendall’s Tau), effectively measure performance and ranking quality against baselines (WACO+FA/FM, no-transfer models). The use of the SuiteSparse Matrix Collection (715 real-world matrices) as a benchmark dataset is appropriate, offering a diverse, established standard for sparse tensor research. Testing on SPADE and A100 GPU aligns with the focus on emerging and established accelerators. Overall, the methods and criteria are sensible and relevant to the application.", "theoretical_claims": "The submission primarily focuses on empirical contributions rather than theoretical proofs, so there are no formal mathematical proofs to check for correctness.", "experimental_designs_or_analyses": "Dataset, Training, and Evaluation Setup (4.1):\n- Design: Uses SuiteSparse Matrix Collection (1500 matrices for training, 715 for evaluation) across CPU, SPADE, and A100 GPU, with 100 random configurations per matrix.\n- Assessment: Sound dataset choice; widely accepted benchmark. Random sampling ensures diversity. Separation of training and evaluation sets avoids overlap. No major issues.\n\nTransferability to SPADE (4.2):\n- Design: Compares COGNATE (Top-1/Top-5) speedups vs. baselines (zero-shot, no-transfer, WACO+FA/FM) using geomean speedups.\n- Assessment: Valid comparison with clear metrics. Optimal speedup baseline (exhaustive testing) is a strong reference. Fine-tuning with 5 matrices is justified later (4.4). No significant flaws.\n\nTransferability to GPU (4.3):\n- Design: Extends evaluation to A100 GPU, comparing COGNATE to cuSPARSE and modified WACO models.\n- Assessment: Logical extension to test generalizability. Consistent methodology with SPADE. cuSPARSE as a baseline is relevant. No issues.", "relation_to_broader_scientific_literature": "- Segregation and Encoding of Program Configurations:\nBuilds on transfer learning concepts from Weiss et al. (2016) and Zhuang et al. (2020), extending homogeneous transfer (e.g., Sasaki et al., 2022) to heterogeneous CPU-to-accelerator settings.\nFeature reuse and latent encoding draw from Neyshabur et al. (2020), adapting them to sparse tensor optimization, unlike prior feature augmentation (Daumé III, 2009; Duan et al., 2012) that struggled with sparsity.\n\n- Data-Frugal Cost Model Framework:\nEnhances WACO (Won et al., 2023) by refining its SCNN-based architecture (Graham & Van der Maaten, 2017) for transferability, contrasting with data-intensive models like Ansor (Zheng et al., 2020).\nAligns with few-shot learning ideas (Shen et al., 2021), reducing data needs compared to traditional ML-based optimization (Chen et al., 2018b; Baghdadi et al., 2021).\n\n- Evaluation on Sparse Accelerators:\nExtends sparse tensor optimization from CPU/GPU systems (Kjolstad et al., 2017; Ye et al., 2023) to emerging accelerators like SPADE (Gerogiannis et al., 2023), complementing hardware-specific efforts (Hegde et al., 2019; Li et al., 2023).\nSpeedup results (1.47× SpMM, 1.39× SDDMM) improve on prior sparse kernel optimizations (Hong et al., 2019; Jiang et al., 2020), offering a data-driven alternative to analytical models (Jin et al., 2024).\n\nCOGNATE bridges ML-based program optimization and hardware acceleration, advancing sample efficiency and cross-platform applicability beyond prior works.", "essential_references_not_discussed": "- Sparse Accelerator Design Context:\nMissing: \"OuterSPACE\" by Parashar et al. (MICRO 2017) introduced a sparse tensor accelerator with configurable tiling, relevant to COGNATE’s mapping of optimizations like tiling and barriers.\n\n- Transfer Learning Efficiency:\nMissing: \"MetaTune\" by Lee et al. (MLSys 2023) proposed a meta-learning approach for tuning tensor programs across platforms with minimal data, akin to COGNATE’s few-shot fine-tuning.\n\n- Cost Model Accuracy:\nMissing: \"AutoTVM\" by Chen et al. (OSDI 2018) (beyond the cited TVM paper) detailed a tuning framework with cost models for sparse workloads, achieving near-optimal schedules.", "other_strengths_and_weaknesses": "Strengths:\n- Originality: Creatively combines transfer learning, feature segregation, and latent encoding to address sparse tensor optimization on emerging accelerators, a novel synthesis of ideas from Neyshabur et al. (2020) and Won et al. (2023).\n- Significance: Tackles a real-world bottleneck in early-stage accelerator design, offering a practical, data-frugal solution with significant speedups (up to 5.46×), impactful for hardware-software co-design.\n- Clarity: Well-structured with clear figures (e.g., Figure 3) and detailed appendices, making complex concepts accessible.\n\nWeaknesses:\n- Originality: While innovative, it heavily builds on WACO (Won et al., 2023), potentially limiting its perceived novelty in the ML optimization space.\n- Significance: Generalizability to untested accelerators (e.g., PIUMA, Vesper) is speculative, tempering its broader impact claims.", "other_comments_or_suggestions": "Could use a brief table summarizing key hyperparameters for autoencoders (like Table 3 for main model) to improve reproducibility.", "questions_for_authors": "How was the 5-matrix fine-tuning sample size chosen beyond empirical observation (e.g., statistical power analysis), and how sensitive are results to slight variations (e.g., 3 or 7 matrices)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VwhBpFCmfd", "reviewer": "Reviewer_HuD5", "summary": "This paper introduces COGNATE, a framework designed to optimize sparse tensor programs (e.g., SpMM and SDDMM) for emerging hardware accelerators using transfer learning. The key innovation lies in leveraging inexpensive data from general-purpose hardware (e.g., CPUs) to pre-train cost models and then fine-tuning them with minimal data from target accelerators. COGNATE addresses challenges such as heterogeneous program configuration spaces and high sample efficiency requirements by:\n\n1. Approximate mapping of comparable code optimizations: Identifying homogeneous features across hardware platforms.\n2. Latent encoding of hardware-specific features: Using autoencoders to compress heterogeneous configurations into a unified latent space.\n3. Few-shot fine-tuning: Achieving competitive performance with only 5% of the data required by baseline methods.\n\nExperimental results demonstrate significant speedups: $1.47\\times$ (SpMM) and $1.39\\times$ (SDDMM) on the SPADE accelerator, and $1.17\\times$ (SpMM) and $1.15\\times$ (SDDMM) on an NVIDIA A100 GPU, compared to existing techniques.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for recognizing the significance of the problem we address and the contributions of our work. We are especially grateful for the time and effort you dedicated to providing such a detailed and thoughtful review. We hope that our following response addresses your suggestions and questions.\n\n> ***Weakness 01 & Question 01*** \\\n> The types of sparse operations supported in this paper are somewhat limited, covering only SpMM and SDDMM. … Have you conducted experiments on a broader range of sparse operations, such as SpMV and MTTKRP?\n\nWe thank the reviewer for this important question. We agree that our current scope considers only SpMM and SDDMM sparse operations. This limitation stems from the fact that these are the only sparse operations currently natively supported by the SPADE accelerator [1] and the SparseTIR framework [2]. \n\nThat said, it is possible to indirectly express SpMV in SPADE using a workaround if the operation is expressed as a SpMM with a very skinny zero-padded dense matrix. During our data collection process, we had gathered performance numbers for a split factor of 16 (an SpMM with 256 dense columns is broken down into 16 SpMMs, each with 16 dense columns). Each of these smaller SpMMs is computationally equivalent to a zero-padded SPADE SpMV. Since we had these performance numbers already available, we trained a model for SpMV during the rebuttal period to evaluate COGNATE’s ability to generalize. The results were promising, with COGNATE achieving a 1.25× geometric mean speedup (top-1) prediction, compared to the optimal speedup of 1.36×. These findings demonstrate that COGNATE could potentially generalize to SpMV-style workloads. We will include these findings and the above clarification in the final version of the paper.\n\n> ***Other Comments*** \\\n> ... Have you considered applying the autoencoder to transform all features, rather than only the heterogeneous ones? If so, what were the results of this approach?\n\nWe thank the reviewer for this insightful question. We share your concern about the need to have mapping functions as diverse new hardware platforms emerge. However, our results suggest that the inclusion of these enabled us to overcome the limitations of prior work in the domain, which achieved effective knowledge transfer only between hardware platforms of the same architecture.\n\nWe explored the idea of applying the autoencoder to all features (both homogeneous and heterogeneous), but found that this approach performed poorly in our experiments compared to the solution we propose in COGNATE. Specifically, we observed the following results for SpMM for the SPADE accelerator when using the autoencoder to encode all input features;\n\n- Top-1 speedup: 1.118× (comapred to 1.40× in COGNATE)\n- Top-5 speedup : 1.237× (comapred to 1.47× in COGNATE)\n\nOur results suggest that transforming homogeneous features, which tend to be consistent and interpretable across hardware platforms, via an autoencoder along with the rest of the features could introduce unnecessary complexity, which hinders the model's ability to generalize effectively. By contrast, limiting the autoencoder to heterogeneous features allows us to preserve the generalizability of shared input characteristics. We will include these findings in the final version to clarify our design choices.\n\n> ***Question 02*** \\\n> On GPUs, there are explicit data movement scheduling parameters, such as cache_read and cache_write. How did you handle these scheduling parameters in your approach? ...\n\nWe thank the reviewer for pointing this out. In our current experiments, we followed the default behavior used in SparseTIR examples, where cache_read and cache_write scheduling optimizations are enabled for SDDMM and disabled for SpMM. As a result, we did not explicitly include these parameters as part of the search space in our evaluation. However, if we were to expand the search space to include these scheduling parameters, we agree that they would need to be treated as heterogeneous features, since they represent architecture-specific optimizations primarily exposed in GPU environments. In that case, they would be encoded through the autoencoder, rather than mapped as homogeneous features. We will explore the possibility of expanding the search space to incorporate these optimizations into our experiments in the final version of the paper.\n\n[1] Gerogiannis, Gerasimos, et al. \"Spade: A flexible and scalable accelerator for spmm and sddmm.\" Proceedings of the 50th Annual International Symposium on Computer Architecture. 2023.\n\n[2] Ye, Zihao, et al. \"Sparsetir: Composable abstractions for sparse compilation in deep learning.\" Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3. 2023.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria, including the use of benchmark datasets, are generally sound and appropriate for the problem at hand. I did not identify any significant issues to address.", "theoretical_claims": "In Section 3.2, this paper employs a mapping function to illustrate the correlation between certain scheduling parameters across different types of machines, suggesting that they can be transformed into one another. However, in the discussion of loop strip mining, the explanation regarding the barrier on SPADE and the reorder operation on the CPU lacks clarity.\n\nFrom my understanding, reordering on the CPU can involve multiple parameter combinations, whereas the barrier is a boolean parameter, making direct conversion between them infeasible. \n\nBesides, there is an indexing error in the argument for loop reordering in the next paragraph.", "experimental_designs_or_analyses": "Since there is currently no cost model that leverages transfer learning to predict the performance of sparse tensor programs across different types of machines, this paper primarily compares the program performance obtained by various methods against the baseline.\n\nFrom an experimental perspective, a comprehensive set of evaluations has been conducted, mainly including:\n\n- Performance achieved by different transfer learning methods (Figure 4)\n- Performance improvements contributed by individual components (Figure 7)\n- Impact of different component choices on performance (Figure 9)\n- Experiments on the data cost associated with transfer learning (Figures 10 and 12)", "relation_to_broader_scientific_literature": "The work builds on: \n\n- ML-based cost models (WACO, TIR) for sparse tensor optimizations.\n- Transfer learning techniques (e.g., feature augmentation/mapping) but classifies machine-specific features into homogeneous and heterogeneous categories for heterogeneous hardware. \n- Hardware-aware optimizations (SPADE, HotTiles) by integrating learned models into accelerator design.\n \n The novelty lies in bridging sparse tensor program optimization and heterogeneous transfer learning, addressing the gap in early-stage accelerator development.", "essential_references_not_discussed": "As far as I know, there aren't Essential References Not Discussed.", "other_strengths_and_weaknesses": "Weaknesses: \n\n- Strengths:\n  - Clarity: Well-structured  with clear methodology.\n- Weaknesses:\n  - The types of sparse operations supported in this paper are somewhat limited, covering only SpMM and SDDMM. So the applicability to other sparse operations (e.g., convolution) is unclear. In contrast, WACO, which this paper compares against, also includes SpMV and MTTKRP.", "other_comments_or_suggestions": "Since feature mapping requires manually identifying similar types of scheduling parameters and defining the mapping function, each time a new hardware platform emerges, a new mapping function needs to be designed. In contrast, the autoencoder-based latent representation proposed in this paper can automatically extract key information from the features. Have you considered applying the autoencoder to transform all features, rather than only the heterogeneous ones? If so, what were the results of this approach?", "questions_for_authors": "1. Have you conducted experiments on a broader range of sparse operations, such as SpMV and MTTKRP? \n2. On GPUs, there are explicit data movement scheduling parameters, such as cache_read and cache_write. How did you handle these scheduling parameters in your approach? Did you treat them as homogeneous features that can be mapped across different hardware, or were they categorized as heterogeneous features and learned through the autoencoder?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5tfAijNBhi", "reviewer": "Reviewer_G9Vs", "summary": "The paper proposes a framework that trains a cost model for performance prediction on general-purpose hardware and then performs few-shot fine-tuning on emerging hardware accelerators. It focuses on optimizing sparse tensor programs on hardware accelerators. The proposed method achieves better hardware performance compared to other approaches like zero-shot and no-transfer learning. For instance, the source model is trained with data from 100 matrices, while the few-shot fine-tuning uses data from only 5 matrices. The paper claims that its few-shot learning approach can fine-tune the cost model with near-optimal accuracy using significantly smaller samples. These claims are generally supported by experimental data.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for recognizing the significance of the problem we address and the contributions of our work. We are especially grateful for the time and effort you dedicated to providing such a detailed and thoughtful review. We hope that our following response addresses your suggestions and questions.\n\n> ***Questions 01*** \\\n> When the paper states, \"The optimal speedup was determined by running all possible configurations (Section 4.2),\" how many configurations are there?\n\nThank you for pointing this out. We apologize for this incomplete statement. We will correct this to be read as:\n\n> “The optimal speedup was determined by running all possible configurations **within our defined search space**.”\n\nand, we will also provide more explanation of how the search space was constructed.\n\nTo clarify, the search space we considered for the SPADE accelerator consists of a total of 256 configurations, derived from combinations of key tunable parameters related to tiling, barrier, cache bypassing, and matrix reordering.\n\nAs shown in Table 1, the parameters for barrier, cache bypassing, and matrix reordering are binary (enabled or disabled). For tiling, we had three numerical parameters: row panels, column panels, and split factor. We had to select a constrained set of values for those numerical parameters since testing every single value would explode the size of the search space and would make data collection infeasible, as explained later.  We spaced those values to resemble the ones tested in the SPADE paper, as those are expected to show more significant performance deviations for different sparse matrices. \n\nIn summary, our defined search space included 4 values for row panels [4, 32, 256, 2048], 4 values for column panels [1024, 16384, 65536, NUM_MATRIX_COLS] (where NUM_MATRIX_COLS depends on each matrix), and 2 values for the split factor [32, 256]. Combined with the three binary parameters, this resulted in 4 × 4 × 2 × 2 × 2 × 2 = 256 configurations.\n\nWe opted for a constrained search space to make data collection feasible. Even so, collecting performance data for training, validation, and testing required approximately 4 million CPU hours. Although we parallelized the required experiments across 32 machines, each with 64 CPU cores, data collection for the constrained set still took nearly three months. Hence, data collection for the full theoretical search space would be infeasible. This challenge again reinforces the importance of data-efficient solutions like COGNATE.\n\nOnce again, we thank the reviewer for raising this point, and we will include this clarification in the final version of the paper to provide additional context around our experimental setup.\n\n> ***Weakness 01*** \\\n> The clarity of some sections, particularly the equations in the code optimization section (Section 3.2), is lacking. Clarifying the notation in Section 3.2 would greatly improve readability.\n\nWe thank the reviewer for highlighting this valuable point. We will add more explanations in the final version. Further, we will include illustrative examples in the appendix to demonstrate how the equations are applied in practice. We believe these additions will make the technical content more accessible and further strengthen the presentation of our contributions. \n\nFor instance, the following data point demonstrates how a SpMM schedule configuration in SPADE can be approximately mapped into its corresponding low-level loop configurations. The original high-level configuration specifies the key tunable parameters used in SPADE:\n\n> name, row_panels, column_panels, split,barrier, bypass, reorder, time \\\n> 144, 4, 1024, 1, 0, 0, 0, 38.83592\n\nThese values are mapped into the corresponding loop-level configuration by applying the equations defined in our paper. Specifically, row_panels, column_panels, and split are used to derive i_split, j_split, and k_split, which represent how the loop indices are partitioned across dimensions. The loop nest structure is encoded by loop_1 through loop_7, which define the execution order of the tiled loops. The binary flags barrier, bypass, and reorder are retained to reflect platform-specific scheduling configurations. The resulting mapped configuration looks like this:\n\n> name, i_split, j_split, k_split, loop_1, loop_2, loop_3, loop_4, loop_5, loop_6, loop_7, barrier, bypass, reorder, time \\\n> 144, 4, 1024, 32, 6, 7, 2, 4, 1, 3, 5, 0, 0, 0, 38.83592\n\nThis representation captures tiling structure, loop ordering, and other applicable optimizations. We will include similar examples along with the corresponding equations and code segments in the final version of the appendix.\n\n> ***Weakness 02*** \\\n>  Some figures (e.g., Figures 5 and 6) are too small to read, which hinders the comprehension of the results.\n\nThank you for pointing this out. We will update these figures to ensure all elements are legible and better support the interpretation of the results in the final version of the paper.", "claims_and_evidence": "The paper asserts that its few-shot fine-tuning approach leads to near-optimal accuracy with significantly fewer samples. These claims are supported by empirical results, including performance comparisons against zero-shot and no-transfer approaches. The evaluation is rigorous and shows clear improvements in hardware performance across various platforms.", "methods_and_evaluation_criteria": "The evaluation focuses on SpMM and SDDMM computations targeting three platforms: CPU, GPU, and an accelerator for sparse matrices. The baseline is adequate, and the dataset consists of 715 real-world matrices, which is sufficiently large. The proposed framework is compared with WacoNet, making the comparison fair.", "theoretical_claims": "The paper does not make any theoretical claims.", "experimental_designs_or_analyses": "The experimental design includes adequate baselines and comparisons against different approaches based on the geometric mean speedup. The study also features good ablation experiments to show the impact of excluding individual components.", "relation_to_broader_scientific_literature": "The paper adequately cites state-of-the-art literature, including WacoNet, which was the previous SOTA method.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n* The paper presents an effective approach that improves upon WacoNet, which was the SOTA.\n* It introduces a latent encoder to capture heterogeneous components and enhances the configuration mapper to optimize techniques unique to individual platforms, improving prediction accuracy.\n* The empirical results demonstrate significant improvements in hardware performance.\n\nWeaknesses:\n* The clarity of some sections, particularly the equations in the code optimization section (Section 3.2), is lacking. The notations are not clearly explained and are difficult to follow.\n* Some figures (e.g., Figures 5 and 6) are too small to read, which hinders the comprehension of the results.", "other_comments_or_suggestions": "* Clarifying the notation in Section 3.2 would greatly improve readability.\n* Increasing the size of Figures 5 and 6 would help readers better understand the visualized data.", "questions_for_authors": "When the paper states, \"The optimal speedup was determined by running all possible configurations (Section 4.2),\" how many configurations are there?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The authors generally found the methodology clear, with clear novelty differentiating it to comparable the closets existing approaches in the literature, and appropriate baselines of comparison. Significant speedups are demonstrated on both the SPADE and A100 GPU accelerators, the method is data efficient, and the paper relatively well written. 

While three of the reviewers are in favour of acceptance of the work post-rebuttal, one reviewer advocates rejection. The primary concern of the reviewers on the negative side is based in the main motivation of the proposed methodology: the method's ability to generalize to different/emerging hardware accelerators. Specifically, two of the reviewers' concerns centred around the evaluation on only two hardware platforms, which does appear to be a significant limitation on the author's claim of generalizability. While the authors address this in the rebuttal as a lack of availability of other hardware architectures, reviewers did not find this convincing, pointing out there are other platforms available to researchers or even simply different GPU architectures that would allow the authors to better demonstrate this strong claim holds.
Decision: Accept (poster)