Abstract
The explosive growth of machine learning applications has consequently created a demand for high-performance implementations of tensor contractions, both for dense and sparse tensors. Compilers, code generators and libraries are often limited in what sparse tensor representations are supported. We observe that tensor contractions can be thought of as iterating over the elements of a sparse tensor to perform an operation and accumulation; co-iteration over multiple tensors can be implemented with iteration and lookup. We recognize that the resulting code can be restructured by specifying a computation, its data layout, and how to iterate over that. We illustrate the need for this generality with two different implementations of block-based data layouts implementing sparse matrix-vector multiplication (SpMV). We show how to generate these implementations with a tensor iterator abstraction designed to be integrated into the MLIR compiler, and present measurements of nearby manual implementations to demonstrate the tradeoffs and complexities with these different implementations.
Supported by the DOE Exascale Computing Project.
This is a preview of subscription content, log in via an institution to check access.
Access this chapter
Tax calculation will be finalised at checkout
Purchases are for personal use only
References
Bik, A., Koanantakool, P., Shpeisman, T., Vasilache, N., Zheng, B., Kjolstad, F.: Compiler support for sparse tensor computations in mlir. ACM Trans. Archit. Code Optim. 19(4) (2022). https://doi.org/10.1145/3544559
Bik, A., Koanantakool, P., Shpeisman, T., Vasilache, N., Zheng, B., Kjolstad, F.: Compiler support for sparse tensor computations in MLIR. ACM Trans. Arch. Code Optim. 19(4), 1–25 (2022). https://doi.org/10.1145/3544559
Bik, A.J.C.: Compiler Support for Sparse Matrix Computations. Ph.D dissertation, Leiden University (1996)
Buluç, A., Williams, S., Oliker, L., Demmel, J.: Reduced-bandwidth multithreaded algorithms for sparse matrix-vector multiplication. In: 2011 IEEE International Parallel & Distributed Processing Symposium, pp. 721–733 (2011). https://doi.org/10.1109/IPDPS.2011.73
Evtushenko, G.: December 2019. https://medium.com/gpgpu/block-sparse-matrix-vector-multiplication-with-cuda-4e616b30267
Kjolstad, F., Kamil, S., Chou, S., Lugato, D., Amarasinghe, S.: The tensor algebra compiler. Proc. ACM Program. Lang. 1(OOPSLA) (2017). https://doi.org/10.1145/3133901
Kotlyar, V., Pingali, K., Stodghill, P.: A relational approach to the compilation of sparse matrix programs. In: Lengauer, C., Griebl, M., Gorlatch, S. (eds.) Euro-Par 1997. LNCS, vol. 1300, pp. 318–327. Springer, Heidelberg (1997). https://doi.org/10.1007/BFb0002751
Lattner, C., et al.: Mlir: scaling compiler infrastructure for domain specific computation. In: 2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO), pp. 2–14. IEEE (2021)
Zhao, T., Popoola, T., Hall, M., Olschanowsky, C., Strout, M.M.: Polyhedral specification and code generation of sparse tensor contraction with co-iteration (2022). https://doi.org/10.48550/ARXIV.2208.11858, accepted for publication, ACM Transactions on Architecture and Code Optimization
Acknowledgments
This research was supported in part by the Exascale Computing Project (17-SC-20-SC), a joint project of the U.S. Department of Energy’s Office of Science and National Nuclear Security Administration and by the National Science Foundation under project CCF-2107556.
Author information
Authors and Affiliations
Corresponding author
Editor information
Editors and Affiliations
Rights and permissions
Copyright information
© 2023 The Author(s), under exclusive license to Springer Nature Switzerland AG
About this paper
Cite this paper
Jolly, J., Goyal, P., Sahoo, V., Johansen, H., Hall, M. (2023). Tensor Iterators for Flexible High-Performance Tensor Computation. In: Mendis, C., Rauchwerger, L. (eds) Languages and Compilers for Parallel Computing. LCPC 2022. Lecture Notes in Computer Science, vol 13829. Springer, Cham. https://doi.org/10.1007/978-3-031-31445-2_2
Download citation
DOI: https://doi.org/10.1007/978-3-031-31445-2_2
Published:
Publisher Name: Springer, Cham
Print ISBN: 978-3-031-31444-5
Online ISBN: 978-3-031-31445-2
eBook Packages: Computer ScienceComputer Science (R0)