# Scalable Block-Sparse Matrix Multiplication Using
Template Task Graphs

Joseph Schuchart, Aurelien Bouteiller, Thomas Herault, Edvard Valeev,
George Bosilca, Robert J Harrison

► To cite this version:

Joseph Schuchart, Aurelien Bouteiller, Thomas Herault, Edvard Valeev, George Bosilca, et al. Scalable Block-Sparse Matrix Multiplication Using Template Task Graphs. WAMTA 2025 - Workshop on Asynchronous Many-Task Systems and Applications, Feb 2025, Saint-Louis, Missouri, United States. pp.120-132, 10.1007/978-3-031-97196-9_10. hal-05423357

HAL Id: hal-05423357

https://inria.hal.science/hal-05423357v1

Submitted on 18 Dec 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
# Scalable Block-Sparse Matrix Multiplication Using Template Task Graphs

Joseph Schuchart¹[0000-0003-2041-7877], Aurelien

Bouteiller²[0000-0001-5108-509X], Thomas Herault²[0000-0001-6756-6189], Edvard

Valeev³[0000-0001-9923-6256], George Bosilca²[0000-0003-2411-8495], and Robert

J. Harrison¹[0000-0002-8777-7466]

¹ Stony Brook University, Institute for Advanced Computational Science, Stony Brook, NY, USA

² Innovative Computing Laboratory, The University of Tennessee, Knoxville, TN, USA

³ Virginia Tech, Blacksburg, VA, USA

**Abstract.** Block-sparse matrix operations are a special case of general sparse algebra where the matrix is sparsely populated with dense blocks, e.g., in sparse tensor algebra for quantum chemistry. One of the challenges of implementing distributed matrix multiplication $C = A \times B$ in general is the management of communication flows since both input matrices $A$ and $B$ are readily available and must be distributed to the processes computing the relevant blocks of $C$. In this paper, we propose an addition to the Template Task Graph programming model that allows applications to constrain the execution of tasks using a flexible API. We show that such constraints can be used in a pure dataflow model to replace artificial control flow with a more structured approach. In the context of sparse matrix multiplication, we found that constraints allow us to limit the number of concurrent communications and thus avoid creating a bottleneck in the network.

**Keywords:** Template Task Graph · Sparse Matrix Multiplication · Scheduling Constraints

# 1 Introduction

Block-sparse matrix operations are a special case of general sparse algebra where the matrix is sparsely populated with dense blocks. These operations are used in various applications, e.g., for sparse tensor algebra in quantum chemistry [4,6].

In this work, we focus on the an implementation of block-sparse matrix multiplication, with $C = A \times B$ computing the product of a matrix $A$ times a matrix $B$ and storing the result in a matrix $C$. This forms the basic operation for many of the applications mentioned above. We assume that all matrices are distributed across a $NP = P \times Q$ process grid and $A$ consisting of $M \times K$ rows and columns in $NT \times KT$ blocks, $B$ consisting of $K \times N$ elements in $KT \times NT$ blocks, and $C$ consisting of $M \times N$ elements in $MT \times NT$ blocks. Matrix blocks are dynamically sized based on the structure of the $A$ and $B$ matrices, generally in the range
[t, T]. Blocks are distributed across processes in 2D block-cyclic distribution, as shown in Figure 1.

<figure><img src="image_2.png" /><figcaption>Fig. 1: Block-sparse matrix multiplication with 2D block-cyclic distribution visualized by coloring all blocks that belong to the same process with the same color. The corresponding MPI rank is also shown on some blocks. Shaded blocks are zero, solid blocks are non-zero.</figcaption></figure>

For a dense matrix, the result for a block $C_{i,j}$ is the sum of the products of a row $i$ of A and a column $j$ of B, i.e.,

$$C_{i,j} = \sum_{k=0}^{KT} A_{i,k} B_{k,j}.$$

This can be extended for sparse matrices to limit the set of operations to be performed to operations involving only non-zero blocks, i.e.,

$$C_{i,j} = \sum_{k=0}^{KT} A_{i,k} B_{k,j} \forall A_{i,k} \neq 0 \land B_{i,k} \neq 0.$$

The main challenge of implementing distributed matrix multiplication in a task-based programming model are two-fold. First, given a tile distribution of A and B, their tiles must be communicated to other processes such that all $A_{i,k}$ are matched with the respective $B_{k,j}$ to ensure that all contributions to $C_{i,j}$ are computed. Simply sending out all tiles of A to the respective owners of B will likely cause a spike in network traffic that will lead to bottlenecks. It is thus desirable to space out these communications.

Second, provided that the respective $A_{i,*}$ and $B_{*,j}$ are available at the respective processors, all $NT^2$ computations of $C_{i,j}$ for a given $k$ are eligible for compute. This abundance of concurrency may lead to undesirable resource contention, e.g., thrashing of data on an accelerator if not all $C_{i,j}$ computed by the
local process fit onto the device. It also may lead to the instantiation of $NT^3$ tasks, of which only $NT^2$ are eligible to execute at a time.

In this work, we will focus on the first challenge, i.e., the network contention caused by the initial transfer of tiles of A and B and how to implement flow-control for these transfers. We will outline how the proposed approach can be applied to controling the concurrency of operation on C.

In this work, we analyze an implementation of the block-sparse matrix-matrix multiplication in the Template Task Graph (TTG) programming model[16], with a focus on the communication requirements. In the TTG model, a data flow-based algorithm is expressed through an abstract graph that is unrolled at runtime based on the flow of data through the template task graph.

# 2 Background and Related Work

## 2.1 Template Task Graph

Task-based programming paradigms have been extensively developed to manage task dependencies in parallel programming. These approaches can be classified based on how the DAG of tasks is represented and managed.

The Dynamic Task Discovery (DTD) paradigm defines the DAG implicitly through the sequential creation and runtime discovery of tasks and their dependencies. Frameworks like OpenMP [2], StarPU [1], OmpSs [5], Legion [3], and PaRSEC DTD [11] rely on this approach. A key challenge with DTD is its limited scalability in distributed-memory systems, as the sequential discovery process must be replicated across all nodes to ensure consistency, leading to inefficiencies at scale.

Frameworks such as C++ Futures [15], HPX [12], MADNESS [9], Charm++ [13], and UPC++ [18] use continuations to explicitly express task dependencies. This approach requires instantiating all edges in the DAG programmatically, which can lead to scalability challenges due to the increased memory and computational overhead of managing a large number of dependencies explicitly. While this model provides a high degree of expressiveness and asynchronous execution, its scalability can be limited for fine-grained tasks.

Frameworks like PaRSEC PTG [7], Dagger [8], and Intel TBB [14] employ an abstract representation of the task graph. These models describe tasks and dependencies explicitly in a declarative format, enabling efficient scheduling and execution while reducing runtime overhead. By decoupling the representation of the DAG from the execution flow, these systems offer scalability and performance benefits, particularly for distributed-memory systems, where they can precompute and optimize task execution plans.

The Template Task Graph<fnref n="4" /> (TTG) paradigm belongs to the third category. It integrates compile-time task graph generation with runtime flexibility, enabling scalable and efficient execution. TTG’s compact and expressive representation balances performance and usability, making it a robust approach for modern parallel programming challenges.

<footnote n="4">⁴ https://github.com/TESSEorg/ttg</footnote>
## 2.2 Distributed Matrix Multiplication

Existing approaches like SUMMA [17] and DBCSR [4] employ a regular communication scheme where rows and columns of A and B are communicated to peers based on k, following the owner-compute principle based on the blocks of the matrix C. In [10], the authors present a parallel implementation of block-sparse matrix multiplication using the PaRSEC task-based runtime system. The implementation uses a two-level broadcast hierarchy to distribute the tiles of A and B to the respective processes, introduce control flow to limitate the memory usage, and rely on a greedy algorithm to schedule the tasks belonging to the current active set of non-zero tiles. DB-SpGEMM [19] extend the work of [4] by using 3D matrix partitioning, a computational task decomposition strategy based on the characteristics of block-sparse matrices, asynchronous task scheduling and multiple optimizations at the communication level to avoid explicit data transposition.

# 3 Block-Sparse Matrix Multiplication in TTG

We had previously presented preliminary results for a block-sparse matrix multiplication [16]. We have since revised the implementation with improved communication management and the constraints that are presented in this paper.

The structure of the BSPMM TTG is shown in Figure 2. We implemented our SpMM as a single class with subclasses representing the different operations involved. At the center of our implementation is the `MultiplyAdd` task, which takes three blocks as input (A, B, and C) and outputs the updated C. The result may flow either back into the next GEMM task for this block (for $k < N$) or flow to a reduction task, which then forwards the block to its final location (`write_spmatrix`).

The structure of the sparse A and B is globally known, i.e., each process can compute a schedule of operations for each $C_i, j$. This schedule is constructed when the 'SpMM' class is constructed and is thus accessible by all operations. The tiles of A, B, and C are block-cyclically distributed across all processes.

The inputs for A and B are read from the input matrices by reader tasks (`read_spmatrix` and sent to broadcast tasks. We chose to implement a two-level broadcast hierarchy: a global broadcast (`bcast_a` and `bcast_b`) and a local broadcast (`local_bcast_a` and `local_bcast_b`). The global broadcast iterates over the set of operations required for the SPMM operation and sends the tiles to these processes. The local broadcast receives the tiles of A and B and forwards the tiles to the respective inputs of the local GEMM operations. The motivation for this two-tier broadcast is to avoid sending the set of locally-known keys along with the tile to each process; on a large-scale run, that list of keys can grow linearly with the problem dimension, leading to a significant overhead in the communication. Since each process knows about its local operations it is possible to generate the list of keys locally and avoid this overhead. It remains as future work to alleviate the requirement of sending key lists in a combined broadcast by using key generators that can be instantiated on each process independently.

Once constructed, each process starts the sparse matrix multiplication by starting the `Control` task locally. For every tile of A and B owned by the local process, the `Control` task instantiates a `read_spmatrix` task. This task receives the tuple `(i,k (A))` or `(k,j (B))` from the `Control` task identifying the tile to read from the input matrix
<figure><img src="image_3.png" /><figcaption>Fig. 2: Template Task Graph of block-sparse matrix multiply (BSpMM).</figcaption></figure>

and to pass on to the respective broadcast task (`bcast_a` or `bcast_b`). These broadcast tasks will send the respective tile to all processes that require it, based on the prepopulated schedule. The tiles are sent to the `local_bcast_*` tasks on each receiving process, which then consults the prepopulated schedule to distribute the tile to all local `MultiplyAdd` tasks. A `MultiplyAdd` task becomes eligible for execution once all its inputs are available, including the recursion of the tile in *C* from earlier *k*. Here we instantiate up to $NT^3$ tasks, with only the first tasks with *k* = 0 eligible for execution at first. Tasks with *k* > 0 will become eligible once the respective tasks with *k'* = *k* - 1 have completed.

# 4 Existing Approaches in TTG

## 4.1 Priorities

Priorities are a weak form of influencing decisions by the scheduler. In TTG, applications can assign a priority *p* to each task that is treated as a hint by the runtime on the
order in which eligible tasks are to be executed. A thread looking for the next eligible task will make a best effort to find a task with $max(p_i)$ for each eligible task $t_i$ known to the scheduler. However, some schedulers may consider a limited subset of tasks for this search, e.g., in a hierarchical scheduler a thread may only look at the tasks in it's bounded local queue, which may have lower priority than tasks in the queues of threads or the global queue. An exhaustive search for $max(p_i)$ would be cost-prohibitive.

Given the non-binding nature and the fact that priorities only influence the order of execution, they are not suitable to influence or limit the execution of broadcast tasks and thus do little to throttle the injection of data into the network by the broadcast tasks.

## 4.2 Control Edges

Despite its data-flow nature, TTG provides the capability for control flow between tasks through edges that carry a task identifier and void data. Such control flow edges can be used to control the execution of tasks by allowing tasks to be blocked from becoming eligible for execution until its control flow input(s) are triggered by connected tasks. The latter might not otherwise provide data flow inputs to the tasks triggered by the control flow.

For example, in the template task graph depicted in Figure 2 we could introduce control edges to throttle the broadcasts, e.g., from `MultiplyAdd` back to `bcast_b`. This would block the `bcast_b` for any *k* that has not yet been released by `MultiplyAdd` and throttle the distribution of tiles based on which GEMM tasks have been executed. A similar approach could be implemented with a recursive control edge on the broadcast task itself, which would order the broadcasts of tiled based on *k*.

However, such a design is rather error-prone and inflexible. The first case requires careful attention in the `MultiplyAdd` task about which broadcast tasks have yet to be released, and when. In a sparse matrix multiplication, not every *k* may require broadcasts but a missed *k* may lead to broadcast tasks that are stuck and will never be released. Moreover, as we will show later it is beneficial to maintain a sliding window of $k_i...k_{i+w}$ where *w* broadcasts are active at once to maximize the overlap between transfers and compute.

# 5 Scheduling Constraints

The goal of this work was to design a mechanism that allows applications to control the execution of tasks through an abstract interface that does not require manual tracking of outstanding tasks. Such a mechanism should constrain the set of tasks the scheduler can choose while not relying on explicit control flow between tasks. We thus introduce the concept of *scheduling constraints*.

A scheduling constraint is a predicate *P* that is evaluated by the task scheduler to determine whether a task is eligible for execution *after* all of its inputs have been satisfied. The set of eligible, unconstrained tasks $\tau$ with tasks $t_i$ is thus defined as a subset $\tau \subseteq T$ of the set of known tasks for which $P(t_i)$ is true. All other tasks $t_j \notin \tau$ may have their inputs satisfied but are constrained because the predicate $P(t_j)$ is false.

The interaction between the data-flow management determining whether a task is eligible for execution, the constraint, and the priority-aware task execution is depicted in Figure 3. It is important to note that tasks that have their inputs satisfied but
<figure><img src="image_4.png" /><figcaption>Fig. 3: Scheduling constraints are situated between the data-flow management and the priority-aware task execution.</figcaption></figure>

whose execution is constrained by the predicate are managed by the constraint and re-evaluated to ensure their eventual execution, e.g., once execution has advanced and the predicate evaluates to true. In our implementation, it is up to the implementation of the constraint to manage these *deferred* tasks and ensure that they are eventually passed to the runtime scheduler for execution.

## 5.1 Constraints in TTG

In TTG, applications can either use the existing constraint or provide a custom constraint. A constraint is attached to a template task upon creation of the template task and queried by the runtime system when a task first becomes eligible for execution, i.e., all of the task's inputs have been satisfied. The constraint then decides whether the task may be executed or whether it has to be deferred.

It is possible to attach multiple constraints on template tasks and the runtime system will query the constraints in the order in which they were attached, pausing the evaluation of subsequent constraints until the current constraint has released the task. This enables encapsulation and composition of constraints.

An example of a custom constraint is provided in Listing 1.1. The `CountingConstraint` tracks the number of tasks eligible for execution and defers tasks if a maximum count is reached. Each constraint derives from `ttg::ConstraintBase`, which provides an interface for the runtime to register listeners that are triggered when a task is released. The constraint itself must implement two functions: `check()` (Lines 8–13) and `complete()` (Lines 16–23). The `check()` method is called by the runtime once a task becomes eligible for execution the first time to evaluate the predicate. If the predicate evaluates to false, the constraint stores the task identifier (`key`) and a handle to the template task (`tt`, the task template). The latter is needed to distinguish listeners if the constraint is shared by multiple template tasks.

The `check()` method is called once a task is completed to allow the constraint to release further tasks. If deferred tasks exist, the constraint may call `notify_listener()` to notify the runtime system that a new task (or set of tasks) has become eligible for execution.

The constraint is applied to a template task `tt` in Listing 1.2 in Lines 1–4. Subsequently, all instantiations of `tt` will be limited to 42 tasks that are eligible for execution.
```cpp
1 template < Typename Key>
2 struct CountingConstraint : public ttg::ConstraintBase<Key> {
3   CountingConstraint (std::size_t count)
4     : m_count (count) { }
5   /* Evaluate predicate for the given task identified by key. */
6   * If P evaluates to false, store the task for later return. */
7   * Return true if P evaluates to true. */
8   bool check (const Key& key, ttg::TTBase *tt) {
9     if (--m_count > 0) return true; // may execute
10    auto g = this->lock guard ();
11    m_list. emplace_back (key, tt); // store for later release
12    return false;
13  }
14  /* Called by the runtime to signal that the */
15  * given task has completed execution. */
16  void complete (const Key& key, ttg::TTBase *tt) {
17    auto g = this->lock guard ();
18    if (!m_list. empty()) {
19      const auto& [key, tt] = m_list. front (); // pick stored task
20      this->notify_listener (key, tt); // provided by ConstraintBase
21      m_list. pop_front ();
22    } else { ++m_count; } // no task to execute, increment counter
23  }
24  private:
25    std::atomic < std::size_t > m_count;
26    std::list < std::pair < Key, ttg::TTBase*>> m_list;
27};
```

Listing 1.1: A custom constraint that caps the number of eligible tasks and defers tasks if the cap is reached.

```cpp
1 ttg::Edge < int, double > e; // key: int, value: double
2 auto tt = ttg::make.tt (..., ttg::edges (e), ...);
3 // limit execution to 42 tasks at a time
4 tt->add . constraint (CountingConstraint < int> (42));
```

Listing 1.2: Using the counting constraint to restrict tasks of type `tt` to 42 concurrent instances.

at once. Any additional tasks will be stored in the constraint structure's `m_list` and released when a task completes. This can be useful in cases where the system has to process many tasks of type `tt` in addition to other tasks that must be processed simultaneously. By limiting the number of eligible tasks to less than the number of available threads (e.g., 64) the application can ensure that a sufficient number of threads is available for tasks of other types.

## 5.2 Shared Constraints

The constraint discussed above was applied to a single template task. However, there may be cases where one constraint should be applied to multiple template tasks. Simply using priorities would not be sufficient since priorities cannot guarantee a balance between different types of tasks, e.g., a high priority for one task type might drown out another type of tasks. TTG provides a way to create a shared constraint using `ttg::make_shared_constraint()`, which returns a `std::shared_ptr` of the constraint.
This shared pointer can then be passed to `add_constraint()` on multiple template tasks.

While the type of the task identifier (`Key`) is typically fixed, it might be necessary to use different parts of the key to identify an *ordinal* to reliably order keys used in different template tasks. For example, in the example shown in Figure 2 the two template tasks `bcast_a` and `bcast_b` use two-dimensional keys (`Key<2>`) to identify the position of the matrix tile to broadcast. However, the *k* term is in different positions of the key, e.g., `key[1]` for `bcast_a` and `key[0]` for `bcast_b`. Constraints can thus be added to tasks together with a function that maps the key to an ordinal that can be used to compare two keys. If a mapper is provided to `add_constraint()`, the signature selected by the `TT` class to check the constraint and to release a task will include the ordinal.

## 5.3 Sequenced Key Constraint

TTG comes with a predefined constraint, the *sequenced key constraint* implemented in the `ttg::SequencedKeyConstraint` class that we will use for the block-sparse matrix multiplication. Using the *ordinal* described above, the constraint determines a sequence of sets of tasks, i.e., tasks with the same ordinal map to the same set. When the sequence advances, all tasks in that set are released for execution and tasks that are discovered subsequently and would be assigned that set are released immediately.

An example in the context of BSpMM is shown in Listing 1.3. The constraint is added to the broadcast tasks and controlled by the GEMM task, i.e., once all GEMM for a given *k* on the local process have completed the GEMM task releases the next batch of broadcast tasks. This allows us to throttle the communication while allowing for sufficient overlap of communication and computation. The window size is flexible and can be controlled by the application.

In Listing 1.3, a shared constraint is created in Line 14 and subsequently added to the broadcast tasks for tiles in the matrices *A* and *B* in Lines 15 and 16, respectively. Inside the `spmm` task (Lines 18–28), the GEMM kernel is called on the inputs of the task (corresponding tiles from *A*, *B*, and *C*; Line 20).

In order to enable some overlap of compute and communication, we wanted to allow for a *window* of task sets to be active. In Line 22, the code accounts for the completed task. The array `k_cnt[]` contains the number of GEMM tasks the local process has to perform for each *k*. If all GEMM tasks for the current *k* have completed the task determines the next set of broadcast tasks (Line 23) and releases the respective *k* on the constraint (Line 24). The definition of `next_release()` is not shown here but essentially iterates over the `k_cnt` array to find the next non-zero entry with distance of at least `parallel_bcasts`. The array `k_cnt` is populated based on the sparsity structure of *A* and *B* when the BSpMM task is instantiated.

# 6 Evaluation

We evaluate the performance of our approach on Hawk, an HPE Apollo system installed at HLRS in Stuttgart, Germany. Each node comprises two AMD EPYC 7742 CPUs with 64 cores each and 256 GB memory. The nodes are connected through InfiniBand HDR200 in a hypercube topology. All codes were compiled with GCC 10.2 against HPE's Message Passing Toolkit (MPT) 2.26.

We use the sequenced key constraint described in Section 5.3 to limit the number of broadcasts that are active concurrently. The results are shown in Figure 4 where we
1 Matrix A = sparse_matrix(N, K), B = sparse_matrix(K, M);
2 ttg::Edge<Key<3>, Tile> a2g, b2g, g2g; // edges: A->gemm, B->gemm, C->gemm
3 ttg::Edge<Key<2>, Tile> ra, rb; // edges: read matrix -> broadcast
4 // Read tiles of A and broadcast them
5 auto read_a = ttg::make.tt(..., ttg::edges(), ttg::edges ra);
6 auto bcast_a = ttg::make.tt(..., ttg::edges ra), ttg::edges(a2g));
7 // Read tiles of B and broadcast them
8 auto read_b = ttg::make.tt(..., ttg::edges(), ttg::edges rb));
9 auto bcast_b = ttg::make.tt(..., ttg::edges(b2g));
10 // count operations for each k based on A and B
11 std::vector<std::atomic<size_t>> k_cnt = count_k(A, B);
12
13 using constraint_type = ttg::SequencedKeysConstraint<Key<2>>;
14 auto constraint = ttg::make_shared_constraint<constraint_type>(false);
15 bcast_a->add_constraint(constraint, [](const Key<2>& key){ return key[1]; });
16 bcast_b->add_constraint(constraint, [](const Key<2>& key){ return key[0]; });
17
18 auto spmm = ttg::make.tt( [=](const Key<3>& key, const Tile& a,
19                          const Tile& b, Tile&& c) {
20   gemm(a, b, c); // generic GEMM
21   int k = key[2];
22   if (0 == --k_cnt[k]) { // release the next k if possible
23     auto release_k = next_release(k_cnt, k, parallel_bcasts);
24     constraint->release(release_k);
25   }
26   ttg::send<0>(next_k(A, B, key), std::move(c)); // forward C tile
27 }, ttg::edges(a2g, b2g, g2g), // input edges
28   ttg::edges(g2g)); // output edge

<caption>Listing 1.3: Using `SequencedKeyConstraint` in block-sparse matrix multiplication. The local broadcast depicted in Figure 2 are omitted here for clarity.</caption>

compare the performance of different broadcast window sizes against DBCSR under various conditions. We used one process per socket (i.e., two processes per node) to mitigate NUMA effects, each process configured to use 64 threads.

From the results, we can see that for smaller matrices ($N = 128k$, Figure 4a) DBCSR provides higher performance than the TTG implementation. We will discuss this below. TTG, however, achieves higher performance for medium and large matrices ($N = 256k$ and $N = 512k$, Figure 4b-d). In the case of large messages, for (Figure 4c-d) we were unable to run TTG on less than 64 nodes in most configurations due to memory constraints. For the same reason, we were unable to run DBCSR with $NT = 192$ and only starting from 128 nodes with $NT = 256$.

For all sizes, a sizable difference between the various TTG configuration can be observed. In all cases, the use of a single concurrent broadcast column or row yields the lowest performance, although it does enable execution at smaller node counts due to lower memory pressure (Figure 4d). In many cases, the use of 1024 parallel broadcasts yields the highest performance, although it appears that stricter limits can yield higher performance in some cases (e.g., a 10% speedup with 16 concurrent broadcasts at 128 nodes in Figure 4b). Developing a heuristic to automatically choose a value for the number of concurrent broadcasts is left as future work.

We attribute the performance benefits of DBCSR on small matrices to its use of a modified Cannon's algorithm, collective communication operations, and out-of-order updates on C in DBCSR [4], which provides for efficient communication along the row and column of matrices A and B. Integrating collective communication into task-based applications is challenging due to the generally unordered nature in the execution of
<figure><img src="image_5.png" /><figcaption>Fig. 4: Comparison of several broadcast window sizes with DBCSR on Hawk.</figcaption></figure>

tasks. However, constraints could be employed in future work to introduce order into broadcasts and utilize MPI broadcast operations along the rows and columns.

It should also be noted that TTG currently does not support out-of-order execution of the individual GEMM operations on each tile. Even if a process already owns the inputs to a *GEMM(i,j,k)*, this task is delayed until all prior with *k = 0... k - 1* have completed. Introducing out-of-order tasks is left as future work.

# 7 Conclusions

This paper introduced an extension to the template task graph (TTG) programming model that allows applications to control the execution order of tasks by providing constraints that are checked by the runtime system once a task becomes eligible for execution. Upon invocation, constraints may defer the execution of a task with the
ability to later release the task for execution once the requirements of the constraint have been met. We utilize a constraint to control the number of concurrent broadcast tasks in a sparse matrix multiplication benchmark where the results show that limiting concurrency can provide improved performance. Further improvements in TTG are needed to match the performance of DBCSR in all cases. We believe that the inclusion of application-defined scheduling constraints in the scheduling will be useful in resource-constraint executions and other applications contexts, e.g., mixed precision calculations where controlled order of operator application is important to calculate in the lowest possible precision. The application of scheduling to such application domains remains as future work.

## Acknowledgements

This research was supported partly by NSF awards #1931384, #1931387, and #1931347. We gratefully acknowledge the provision of computational resources by the High Performance Computing Center (HLRS) at the University of Stuttgart, Germany.

## References

1. Augonnet, C., Thibault, S., Namyst, R., Wacrenier, P.: StarPU: A unified platform for task scheduling on heterogeneous multicore architectures. Conc. Comp. Pract. Exper. **23** (2011)

2. Ayguadé, E., Copty, N., Duran, A., Hoeflinger, J., Lin, Y., Massaioli, F., Teruel, X., Unnikrishnan, P., Zhang, G.: The design of OpenMP tasks. IEEE Transactions on Parallel and Distributed systems **20**(3), 404–418 (2008)

3. Bauer, M., Treichler, S., Slaughter, E., Aiken, A.: Legion: Expressing locality and independence with logical regions. In: SC '12: Proceedings of the International Conference on High Performance Computing, Networking, Storage and Analysis (2012). https://doi.org/10.1109/SC.2012.71

4. Borštnik, U., VandeVondele, J., Weber, V., Hutter, J.: Sparse matrix multiplication: The distributed block-compressed sparse row library. Parallel Computing **40**(5), 47–58 (2014). https://doi.org/10.1016/j.parco.2014.03.012, https://www.sciencedirect.com/science/article/pii/S0167819114000428

5. Bueno, J., Planas, J., Duran, A., Badia, R.M., Martorell, X., Ayguadé, E., Labarta, J.: Productive programming of GPU clusters with OmpSs. Proceedings of the 2012 IEEE 26th International Parallel and Distributed Processing Symposium, IPDPS 2012 pp. 557–568 (2012). https://doi.org/10.1109/IPDPS.2012.58

6. Calvin, J.A., Valeev, E.F.: TiledArray: A general-purpose scalable block-sparse tensor framework, https://github.com/valeevgroup/tiledarray

7. Danalis, A., Bosilca, G., Bouteiller, A., Herault, T., Dongarra, J.: PTG: An abstraction for unhindered parallelism. Proceedings of WOLFHPC'14 (2014). https://doi.org/10.1109/WOLFHPC.2014.8

8. Gursoy, A., Kale, L.V.: Dagger: combining benefits of synchronous and asynchronous communication styles. In: Proceedings of 8th International Parallel Processing Symposium. pp. 590–596. IEEE (1994)
9. Harrison, R.J., Beylkin, G., Bischoff, F.A., Calvin, J.A., Fann, G.I., Fosso-Tande, J., Galindo, D., Hammond, J.R., Hartman-Baker, R., Hill, J.C., Jia, J., Kottmann, J.S., Ou, M.Y., Ratcliff, L.E., Reuter, M.G., Richie-Halford, A.C., Romero, N.A., Sekino, H., Shelton, W.A., Sundahl, B.E., Thornton, W.S., Valeev, E.F., Vázquez-Mayagoitia, Á., Vence, N., Yokoi, Y.: MADNESS: A multiresolution, adaptive numerical environment for scientific simulation. SIAM J. Sci. Comput. 38(5) (2016)

10. Herault, T., Robert, Y., Bosilca, G., Harrison, R.J., Lewis, C.A., Valeev, E.F., Dongarra, J.J.: Distributed-memory multi-GPU block-sparse tensor contraction for electronic structure. In: 2021 IEEE International Parallel and Distributed Processing Symposium (IPDPS). pp. 537–546 (2021). https://doi.org/10.1109/IPDPS49936.2021.00062

11. Hoque, R., Herault, T., Bosilca, G., Dongarra, J.: Dynamic task discovery in PaR-SEC: A data-flow task-based runtime. In: Proceedings of the 8th Workshop on Latest Advances in Scalable Algorithms for Large-Scale Systems. pp. 1–8 (2017)

12. Kaiser, H., Heller, T., Adelstein-Lelbach, B., Serio, A., Fey, D.: Hpx: A task based programming model in a global address space. In: Proceedings of the 8th international conference on partitioned global address space programming models. pp. 1–11 (2014)

13. Kale, L.V., Krishnan, S.: CHARM++: A portable concurrent object oriented system based on C++. In: Proceedings of the Conference on Object Oriented Programming Systems, Languages and Applications (1993)

14. Kukanov, A., Voss, M.J.: The foundations for scalable multi-core software in intel threading building blocks. Intel Technology Journal 11(4) (2007)

15. Rintala, M.: Handling multiple concurrent exceptions in c++ using futures. In: Advanced Topics in Exception Handling Techniques, pp. 62–80. Springer (2006)

16. Schuchart, J., Nookala, P., Javanmard, M.M., Herault, T., Valeev, E.F., Bosilca, G., Harrison, R.J.: Generalized Flow-Graph Programming Using Template Task-Graphs: Initial Implementation and Assessment. In: 2022 IEEE International Parallel and Distributed Processing Symposium (IPDPS) (2022). https://doi.org/10.1109/IPDPS53621.2022.00086

17. Van De Geijn, R.A., Watts, J.: Summa: scalable universal matrix multiplication algorithm. Concurrency: Practice and Experience 9(4), 255–274 (1997)

18. Zheng, Y., Kamil, A., Driscoll, M.B., Shan, H., Yelick, K.: UPC++: A PGAS Extension for C++. In: 2014 IEEE 28th International Parallel and Distributed Processing Symposium. pp. 1105–1114 (2014). https://doi.org/10.1109/IPDPS.2014.115

19. Zheng, Z., Chen, J., Zhao, Y., Song, L., Qin, X., An, H.: DB-SpGEMM: A massively distributed block-sparse matrix-matrix multiplication for linear-scaling DFT calculations. In: Proceedings of the 53rd International Conference on Parallel Processing. p. 1156–1165. ICPP ’24, Association for Computing Machinery, New York, NY, USA (2024). https://doi.org/10.1145/3673038.3673159