1. Preface​

Development of electronic structure programs aims to solve or discover scientific problems in chemistry and material structure; but the technical challenges involved are often unrelated to chemistry itself. Electronic structure programs need to take into account both 1) development efficiency and reduced communication cost, and 2) program efficiency and resource control; these are also problems that scientific computing in other disciplines, and even general program development tasks, need to solve.

The Rust language has received good reception in some fields of computer science or its applications. But for scientific computing, Rust rarely has outstanding work; and I believe there is still no consensus on whether Rust is suitable for scientific computing.

The previous showcase_rust_riccsd work, I believe, can show that for the problems electronic structure cares about, represented by MP$n$ and CC, it is possible to achieve a good balance between development efficiency and program efficiency with appropriate tools. To demonstrate this possibility, I developed RSTSR as a tensor computing tool. Limited by my horizons, abilities and energy, this tool may not be ideal; but I hope to use this tool to show my understanding of and expectations for Rust scientific computing programs, and to provide some ideas for the development of tensor tools that satisfy electronic structure.

This document hopes to use a simple Q&A style to show my understanding of electronic structure program development and the program problems it cares about, as well as my understanding of the Rust language. The second half of this document will introduce the development ideas of the RSTSR program.

A good program tool should allow users to escape tedious technical details to a certain extent, and focus on their own goals; this is also the original intention of developing math library tools. RSTSR is strongly influenced by NumPy. It is necessary for us to show how NumPy or similar math library tools are concretely used in electronic structure, where their strengths and weaknesses are, and our ideas for improving these problems.

This document assumes that GPU heterogeneity and MPI-scale parallelism are not considered. Some conclusions may not apply to these two situations.

2. Electronic Structure Problems​

2.1. As a Scientific Computing Problem, What Are the Characteristics of Electronic Structure?​

First, I need to state that my understanding of electronic structure is not complete. I have never written CISD, Full-CI (MCSCF), DMRG, or PBC, and I do not understand the scientific computing challenges of these branches.

Only program implementation issues are discussed here, not method development.

Focusing on problems like MP2, CC and DFT, electronic structure programs

1. are mainly composed of matrix operations or operations and contractions of high-dimensional tensors (generally no more than 4-D, rarely exceeding 6-D);
2. have some problems involving eigenvalue solving or matrix decomposition;
3. have a small number of problems involving nonlinear equation solving (iterative solving of matrix equations).

Corresponding to the above 3 problems respectively, some of my views are

1. Arbitrary-dimension tensor operations and contractions can be classified into two kinds of problems: broadcasted matrix multiplication and broadcasted elementwise operations. This will be explained in detail later, but here we only need to know that these two kinds of problems are relatively standard problems, which can be solved with simple programs and BLAS.
2. Eigenvalue problems and matrix decomposition are standard problems, which can be solved with Lapack.
3. This is not a standard problem, and chemists need to design algorithms according to their own needs. Of course, some matrix equation solvers have solutions in Matlab or SciPy; but their efficiency may not satisfy us.

Therefore, among the above 3 problems, the only one where chemists can truly exert their value is the 3rd kind of problem; the rest are standard problems, which can be implemented by program engineers entirely without chemist participation. Chemists are also not professional at dealing with the 3rd kind of problem, and will very likely need to hand it over entirely to numerical mathematicians.

2.2. Why Do Chemists Write Electronic Structure Programs?​

The previous answer denied most of chemists' value in writing electronic structure programs. Taking CCSD as an example, apart from DIIS iteration, chemists seem to have no use; after all, for the remaining problems, one just needs to write the program from the formulas. If DIIS was developed by a numerical mathematician, then chemists would not need to do anything.

But chemists still write electronic structure programs, for the following reasons:

• There is no funding or policy support to hire scientific computing programmers. I might think that because we cannot show society and the public that our work has more important value (which may also be the truth), money and resources cannot reasonably tilt toward us. There is no way around this.
• Method development. Chemists are responsible for improving existing methods. Such improvements are usually about accuracy; they sometimes require chemical intuition, and sometimes are a mathematical structure. New methods may need new program tools; the Davidson diagonalization is believed to have come about this way. Therefore, ideally, chemists also need to be engineers and mathematicians. There are also people who improve efficiency (I think Eshuis of the Furche group improving the RPA algorithm is a typical example), but such work is more like the work of mathematicians, and does not really require chemical participation.

But if it is about implementing existing electronic structure methods more efficiently, or in a new program, or writing the gradient properties of existing electronic structure methods, then these have nothing to do with chemistry.

I have also heard more than one person mention that when working with program engineers, they strongly feel that engineers do not understand electronic structure (which produces a certain degree of negative experience). Since I have never worked with engineers, perhaps my following understanding is wrong: in my view, the possible factors causing this communication barrier are

• either the program engineer is limited in ability and cannot correctly understand numerical problems (unrelated to engineers not understanding chemistry);
• or the chemist did not successfully transform the electronic structure problem into a numerical problem (thus the problem lies with the chemist);

Of course, if one person is both a chemist and an engineer, then this communication barrier does not exist.

2.3. Does the Early Development of the REST Program Need Chemists?​

In my view, no. The main purpose of the early program is to implement existing algorithms (migrating them from other languages or toolchains into REST), which is not chemistry. But an engineer who can understand the language of chemistry is always good.

3. NumPy​

3.1. Why Do You Recommend Chemists to Use NumPy? What Are Its Advantages?​

This is a subjective question. Everyone should have their own opinion on what kind of program tools are suitable for chemists.

The programming needs of chemists should generally fall on method development.

The main reasons I recommend NumPy are as follows:

• Python, the scripting language. Scripting languages usually have an interactive running mode (bash has shell, python has jupyter, etc.), which is an advantage compiled languages do not have (or at least not conveniently).
  But this does not explain why not use Pytorch, Matlab, Mathematica, JavaScript or Julia.
• Arbitrary-dimension tensor support and Basic Slicing. Basic slicing is a means of giving sub-matrices/tensors without copying. This pattern should have started from Fortran, and is not a feature unique to NumPy.
  But this does not explain why not use Pytorch, Matlab, Mathematica, JavaScript or Julia.
• Support for Einsum. This is of great help for prototype implementation of post-HF methods.
  But this does not explain why not use Pytorch, Matlab, Mathematica, JavaScript or Julia.
• Fast matrix multiplication. NumPy connects to high-efficiency BLAS functions for matrix multiplication, so if the computational bottleneck is in matrix multiplication, NumPy is generally not too slow.
  But this does not explain why not use Pytorch, Matlab, Mathematica, JavaScript or Julia.
• Relatively complete linear algebra support. Combined with SciPy, NumPy has solutions for matrix linear algebra, eigenvalue solving, FFT, ODE, extremum problems, etc.
  Finally, there are ODE and extremum problems that JavaScript's stdlib does not support, so it is out.

Among the above 5 factors, the first 3 are decisive factors, all indispensable. But at the same time, you will find that no recommendation reason is unique to NumPy. So it is completely fair to say that recommending NumPy is definitely pushing a personal agenda.

Where NumPy truly has additional advantages over other tools is in

• Python as a general-purpose language. This is something Matlab and Mathematica cannot do, and Julia is not good at. JavaScript is more suitable for the frontend.
• The Python ecosystem. A larger ecosystem means that if you encounter any problem, others may have encountered it too, and there are public solutions. PyPI and Conda are also important to the Python ecosystem.
• Convenience of installation and a smaller runtime. This is relative to Pytorch. Runtime means runtime; here we mean that NumPy as a Python library has a relatively small binary size, and has no extra dependencies besides BLAS.

These additional advantages are not decisive. Moreover, NumPy also has disadvantages; some disadvantages are fatal.

3.2. What Are NumPy's Disadvantages?​

From my usage experience,

1. No automatic differentiation. This is not the most concerning problem for electronic structure programs, but it is indeed one of the most demanded features right now.

2. No heterogeneous support. Using only NumPy, one cannot run on GPU. GPU is indeed one of the important directions of scientific computing in the future.

3. Python loops are slow. This is well known. For many standard matrix operation problems, we indeed may not need Python for loops; but if we encounter a triangular matrix assignment problem like the following (possibly used for tensor symmetrization $t_{ji}^{ba} := t_{ij}^{ab}$):

   for i in range(nocc):
       for j in range(nocc):
           for a in range(nvir):
               for b in range(a):
                   t2[j, i, b, a] = t2[i, j, a, b]
   

   There is no suitable NumPy function for this type of operation, which produces quite serious computational efficiency problems. Although it has solutions (Numba jit), it is always not very convenient.

4. Python iterators are hard to parallelize with threads. But Python parallelism is only suitable at the process level (multiprocessing), not at the thread level. The above is also an example: the above $t_{ji}^{ba} := t_{ij}^{ab}$ computation can be parallelized with a double loop over the $i, j$ indices. Using fairly standard Python parallel libraries, the effect is a complete mess.

5. Numba JIT programming is not native Python. The above two points can be solved with Numba. Numba is not impossible either, but JIT itself takes time, and Numba syntax is not exactly Python; precompiled JIT is also not impossible, but then what is the difference from using the FFI of a compiled language (like C)? JIT programming is not without thresholds: if input and output types are not clearly written, and parallel is not correctly enabled, Numba's acceleration is not good either.

6. Computation on non-contiguous matrices is very slow and not parallelized. A typical problem is the closed-shell RI-MP2 $O(n^4)$ energy summation problem. Simplified to a concrete expression, when matrices $\mathbf{A}$ and $\mathbf{B}$ are both c-contiguous or f-contiguous, the operations $2A - A^T$ or $A \odot B^T$ are outrageously slow.
   Similar problems also appear when calling SciPy. SciPy is usually friendly to f-contiguous matrices; if the input matrix is c-contiguous, it sometimes performs a very inefficient matrix transpose first.

7. Computation on contiguous vectors is not parallelized. NumPy does optimize operations on contiguous vectors (SIMD vectorization), but this is mostly work from 10 years ago. 10 years ago, on personal computers, single-threaded operations could generally saturate memory bandwidth; but times have changed, and for some problems a single thread can no longer keep up with memory bandwidth speed. It can be considered that as long as the problem is not matrix multiplication or eigenvalue problems, NumPy has some room for efficiency improvement.

8. Some function names are quite bad. For example np.ix_; this may be for historical reasons.

9. Memory control is difficult. This is related to Python's variable lifetimes. Rust can avoid this problem.

Some of the above problems are also problems encountered by Matlab and Mathematica. The Julia language can combine the advantages of both scripting languages and compiled languages; from this perspective it is indeed a tool more suitable for scientific computing; but the development trend and professionalism of the language itself worry me somewhat.

3.3. We Are a Rust Program; Why Mention Python and NumPy?​

We are not saying NumPy is good or bad. But we are writing Rust programs and need to use some math libraries. Therefore we need a general understanding of what features math libraries generally have.

From the perspective of scientific computing, NumPy may not be the best choice. But there are reasons to choose NumPy as the baseline for discussing math libraries:

• I am more used to NumPy, and the most popular computational chemistry program in the open-source community right now is PySCF (judging by GitHub stars). Of course there is personal preference involved.
• The Python language that NumPy is based on was not specifically designed for scientific computing. The same is true for C++/Rust in this respect; in contrast are Julia and Fortran. This makes many problems have to be implemented through functions (or operator overloading) rather than through language syntax (e.g., transpose, slicing, matmul, vector solve, inversion).

4. Strategies for Rapidly Developing Electronic Structure Programs​

4.1. Is There a Strategy for Rapidly Developing Electronic Structure Programs? Why Is Rust Not in This Strategy?​

As mentioned before, the decisive reasons I recommend NumPy are

• scripting language,
• arbitrary-dimension matrix or tensor support,
• Einsum support.

Of course NumPy is not the only framework suitable for these three points; but since Rust is not a scripting language (REPL language), I would consider it not suitable for rapid development.

Rust is not without REPL frameworks (the evcxr framework), but in my experience compilation is still rather slow.

4.2. Is There a Strategy for Rapidly Developing Efficient Electronic Structure Programs?​

With the word "efficient" added to this question, the above answer is no longer useful.

The development strategy I have tried in implementing RI-MP2 static polarizability and RI-RHF first-order gradient is 0. Organize the problem into mathematical formulas, avoiding any chemical description as much as possible;

1. Implement it once in NumPy using Einstein summation;
   • This does not mean it can only be implemented with Einstein summation; it is just that for MP$n$/CC tasks, Einstein is absolutely convenient; but there are also other specific problems suitable for other methods. What is meant here is to translate your chemical problem into a program using any framework you find convenient, and execute it correctly. You need to consider what algorithm to use and how to optimize the algorithm, but correctness comes first, and you do not need to consider the actual program efficiency now.
2. In NumPy, split all Einstein summations into (in priority order)
   • matrix multiplication (matmul, matrix multiplication)
   • vector operations (elementwise)
   • operations like summation (reduction)
   • loop unrolling (for loop) Finally obtain a Python program completely without einsum.
3. In compiled languages like Rust/C, translate the above logic one by one.

Except steps 0 and 1, the above process is fairly mechanical. To put it harshly, AI-assisted programming may actually be able to replace steps 2 and 3. But for computational performance, the programs of steps 2 and 3 are also important.

From my experience, after step 1 is successfully completed, step 2 for problems at the level of RI-CCSD or RI-MP2 static polarizability takes about 0.5–5 days, and step 3 takes about 1–10 days; the exact time depends on work state, work environment and concentration, and has nothing to do with the creative intensity at the time. The most difficult part of this process is generally step 1, that is, whether the chemical problem can be correctly written as a program and executed within the math library framework we are used to. Of course, this varies with the problem.

5. The Relationship between Math Libraries and Electronic Structure​

5.1. Why Waste Time Developing Math Libraries? Isn't Our Goal to Solve Chemistry Problems?​

I would ask in return: why do we use Rust? What is the purpose of giving up the nice scripting languages? Does this really solve chemistry problems?

If one refuses to answer this question, then all the rest of this document is meaningless to the reader. Some situations that seem to answer the question but actually refuse to answer include:

• Our program has already decided to feature Rust.
• C++ is the Sekiro of the programming world; Rust is the Genshin Impact of the programming world. Rust is the pinnacle of abstraction.

I estimate that about 50% of computational chemistry software develops its own math library or binds to external math libraries; some software may even develop two sets, making code reading and communication difficult:

• Q-Chem has developed multiple math libraries internally, including libblas, libmathtools, libtensor.
• PySCF mainly uses the external math libraries NumPy and SciPy, and also uses TBLIS; but it has its own programs for some operators, generally defined in np_helper.c.
• Psi4 has developed its own math library libmints/matrix; judging only from the occ part of Psi4's code, these libraries are mainly used as BLAS wrappers;
• MPQC (as a computational chemistry program that has stopped maintenance) cooperated with other groups on MADNESS (which is still actively maintained).

I dare not say the reasons for developing math libraries are self-evident, but developing math libraries is indeed the choice of many electronic structure developers.

But there are also many people, especially engineers leaning toward high-performance computing, who probably choose not to use existing math libraries. They may need a more flexible program writing style, using at most low-level programs like BLAS or FFT, without heavily using high-level interfaces.

The significance of math libraries is somewhat subtle. In my view, its purpose is: balancing development efficiency and program efficiency.

• For the highest development efficiency, one should generally use scripting languages Python, Matlab, Julia;
• For the highest program efficiency, one should generally directly call low-level BLAS and hand-write some functions.

Therefore, the general requirements for math libraries are

• implement important operators to satisfy general numerical computation needs;
• for the implemented operators, the implementation efficiency should approach the ideal limit;
• the code should be simple and intuitive enough, and not easily misused by users.

Even though the significance of math libraries is subtle, I still think it is necessary to use math libraries in computational chemistry programs. Different people have different judgments on this matter; so this statement can be seen as my philosophical proposition, not necessarily rational. In my previous implementation of RI-MP2 static polarizability, my workflow was to first draft it with NumPy, then implement it in C; it took about a week in total. The C part was mainly implemented by calling BLAS. But calling BLAS is error-prone: it is easy to make mistakes in writing matrix dimensions and leading dimensions without the program reporting errors; such code is also not intuitive. Therefore, to write efficient and correct code faster, I am motivated to develop math libraries.

5.2. Implementing a Math Library Is Too Much Work; Can We Design a Math Library with Electronic-Structure-Specific Features?​

I hold a negative position on this question.

The needs that electronic structure has for math libraries already cover the functionality of most math libraries (represented by NumPy). A math library that satisfies electronic structure needs can generally also satisfy the needs of other disciplines.

The specific requirements electronic structure may have for math library features will be discussed later.

For the difficulty of math library development, I think using an unsuitable or inconvenient framework to develop programs is also a time cost, and developing a framework is also a time cost. It is hard for me to evaluate which one takes more time, but the latter indeed gives people the impression of having done nothing, and of not knowing at which specific step the electronic structure method was not implemented. In fact, other developers besides myself, and non-developer classmates, also cannot endorse developing a math library, because it looks very difficult; and with PyTorch being so powerful now, the marginal benefit of doing this seems lower and lower. Only when the shortcomings of NumPy itself are solved, the advantages of Rust as a newly developed compiled language are endorsed, the developed library is satisfactory in both convenience and performance, and the features are relatively complete, can developing a new math library have positive returns; and even so, it cannot directly compete with PyTorch's features and ecosystem.

5.3. Why Not Consider Using PyTorch's Rust Binding?​

PyTorch is almost the best math library framework at present, even without considering its more powerful automatic differentiation. I also deeply respect the author of PyTorch's Rust binding crate tch-rs; this is absolutely not a simple binding, it is itself a very good solution for a Rust math library.

I think the following factors are not necessarily decisive, but combined I tend not to use PyTorch's Rust binding crate tch-rs:

• PyTorch's runtime is too large, and users may have linking problems (CPU-only is somewhat better); it may be difficult to run on some devices; a pure-Rust framework can reduce these dependencies, at least down to only needing BLAS (and CUDA later if GPU support is added). It is better to treat PyTorch as an optional backend rather than binding to this one framework.
• The crate tch-rs is already quite powerful as a libtorch binding and a native Rust interface; but it is still a foreign library, and we cannot fully expect to smoothly add features to it for electronic structure purposes.
• Getting data in and out of tch-rs tensors currently seems possible only by copying. Rust's standard containers for data are either Vec<T> or &[T], and producing or extracting data has no overhead; but tch-rs hands the data completely over to the C++ part, and the Rust part no longer has control over the data, so if you want to perform operations on the data outside tch-rs, there is great overhead.
• tch-rs does not seem to be actively maintained, but it keeps updating PyTorch versions and dependencies to stay usable (or the library may already be in a completed state and needs no updates). The focus of tch-rs developer (Laurent Mazare) is now the candle library, but candle seems even further from electronic structure.
• Following the above two reasons. I think the reason the author later developed candle is that PyTorch's operators are not necessarily up-to-date in the LLM era; but writing new operators in the PyTorch backend is not convenient unless done in C++; to keep up with LLM development in a framework different from PyTorch, a new machine learning framework was needed. Of course, this is not our concern, but it means that using tch-rs sacrifices considerable flexibility.
• It is not impossible on the tch-rs framework, but support for complex floating point numbers is currently basically zero.

5.4. Why Not Consider Using Other Math Libraries in Rust?​

Other Rust math libraries are not unusable either, but I hope to do better. Also refer to RSTSR Features later.

But I also want to say that when I first started with Rust, my mindset was that others wrote math libraries ten years ago, so we might as well use theirs. I held this attitude for a long time, until I could not conveniently use the ndarray framework to achieve a program efficiency satisfactory to me on the RHF first-order gradient problem; and the RHF first-order gradient problem ideally needs three-dimensional tensor support and indexing support close to basic slicing, and rest_tensors is also not very suitable.

I hope that my mindset in developing a Rust math library is not the Not Invented Here effect. Therefore, I want to use this document to fully state my understanding of the relationship between math libraries and electronic structure, to determine what exactly we want, and why I am not satisfied with existing tools.

5.5. What Math Library Features Do Electronic Structure Programs Need?​

We answer this question in reverse. What features do math libraries generally have? Which of them will be used by electronic structure programs?

After the explosion of Python math libraries in 2016–2017 with the popularity of machine learning, starting from 2020 the community formed the Python array API standard. Although this standard is not mandatory, NumPy 2.0 enabled it; and this standard helps us understand the expectations and basic requirements of general users and developers for math libraries.

The main content of the Python array API standard is in the API specification. The main parts are

• Array object: the definition of arbitrary-dimension dense tensors (including 0-D scalars), and basic operations such as addition, subtraction, multiplication and division.

  • CCSD computation needs the double excitation tensor $t_{ij}^{ab}$; various energy and gradient computations need 3c-2e ERIs like $Y_{ia, P}$ or their Cholesky decomposition; MP2 gradient computation often needs to store three of the dimensions of $t_{ij}^{ab}$. Generally speaking, computational chemistry has a rigid need for 3-D tensors, and 4-D tensors are also often needed. Meanwhile, gradients often have three-dimensional tensors like $h_{\mu\nu}^{t} = \langle \mu | t | \nu \rangle$ (where $t$ represents the three components $x, y, z$); although it can indeed be stored as a Vector of three 2-D matrices, storing it directly as a 3-D tensor is generally more convenient.

• Broadcasting: the broadcast computation of tensors. Its rules are quite complex, but the actual applications are two cases:

  • Direct sum or direct product. In MP2 or CCSD energy computation, one encounters the need for a matrix $\Delta_{ab}^{ij}$ (a 2-D matrix only about the indices $a, b$; the indices $i, j$ are generally used in the outer loop), which needs to be computed by:

    d_ab = e_occ[i] + e_occ[j] - e_virt[:, None] - e_virt[None, :]
    

    This is a standard direct sum computation. It uses both broadcasting and the basic indexing trick of adding a dimension (unsqueeze / newaxis).

  • Operations of matrices with different dimensions. This is the original purpose of broadcasting, not used much in chemistry; but for the cases where it can be used, broadcasting is very convenient. For example, handling the multiplication of gradient matrices and density matrices:

    $$\partial_{A_t} E \leftarrow - \sum_{\mu} \sum_{\nu \in A} h_{\mu\nu}^{t} D_{\mu\nu}$$

    Letting the variable slc be the set of basis functions on atom $A$ (we are now discussing programs based on atomic orbitals),

    de[atom, :] -= (deriv_h[:, :, slc] * rdm1[:, slc]).sum(axis=(-1, -2))
    

    The above computation uses elementwise multiplication of a 3-D and a 2-D matrix, which is broadcasting multiplication.

• Creation Functions: tensor creation functions. Empty tensors, zero tensors, identity matrices.

• Data Types: this is not a concrete function, but requires the tensor library to at least handle 8–64 bit integers and unsigned integers, 32–64 bit floating point numbers and complex floating point numbers, and boolean types.

• Element-wise Functions: operation functions. This includes common functions like abs, sin, log, greater, floor, isnan. Actually many functions are not used in chemistry, but they are still needed to handle the needs of methods like Laplace-Transform for these functions. Determining electron occupation numbers often requires magnitude comparisons.

• Indexing: arbitrary-dimension indexing, generally referring to NumPy's basic indexing, but the Python array API standard also requires boolean tensor indexing. Computational chemistry has needs for both basic indexing and boolean tensor indexing (or similar index-list indexing):

  • In the two code cases of broadcasting, the direct sum computation of $\Delta_{ab}^{ij}$ needs newaxis indexing (adding a dimension via indexing); the computation of $\partial_{A_t} E$ is standard indexing that extracts part of the contiguous atomic orbitals.

  • Some special frozen orbital, or CAS, or MOM orbital selections are implemented by indexing with a list of indices. Suppose we want to obtain the [1, 2, 3, 5, 6] orbitals of the system,

    frz_orbs = [1, 2, 3, 5, 6]
    frz_coeff = mo_coeff[:, frz_orbs]
    

• Linear Algebra Functions: matrix multiplication, matrix transpose, tensor contraction, vector inner product. We generally use the first two, and the importance of the first two is self-evident; tensor contraction is implemented by tensordot, which is also a frequently used feature, but at the cost of code readability or convenience it can be replaced by matrix multiplication.

• Manipulation Functions: common functions include reshape for changing shape, permute_dims (transpose) for tensor transposition, and stack and concat for stacking or concatenating matrices. reshape is a very commonly used and important function.

• Searching Functions: argmax/argmin, nonzero, where. We do not use these functions much.

• Set Functions: convert tensors to sets (set, in the programming context). We do not use these functions.

• Sorting Functions: sorting functions. Not used much in electronic structure, but may be used in the future.

• Statistical Functions: max/min, sum, prod, std, etc. We need to use sum often.

In summary, electronic structure has comprehensive needs for math libraries. Except the three categories Searching, Set, Sorting, we need all the rest.

In addition, there are many features that math libraries may not have, but are of important use in computational chemistry:

• Matrix decomposition. This is part of Lapack; its significance is self-evident.

• Special operators or einsum. einsum is not within the scope required by the Python array API standard, but it is of great help to electronic structure. Many DFT computations are inefficient under standard operators; one either needs to implement a powerful einsum, or manually implement these operators (like the current rest_tensors and PySCF).

• Nonlinear matrix equations. This refers to solving $f(\mathbf{x}) - \mathbf{b} = \mathbf{0}$; but note that although $f(\mathbf{x})$ can be written as $\mathbf{A} \mathbf{x}$, due to computational difficulty or excessive storage, $\mathbf{A}$ will not be directly computed, so $\mathbf{x}$ needs to be implemented by solving equations. The typical problem of this kind is the CP-KS equation, and of course also similar Casida equations or wave function stability analysis. These problems may be problems that SciPy cares about, or may require us to implement the algorithms ourselves.

6. Understanding of Performance, Understanding of Rust​

6.1. How Do You View Rust?​

Rust is a new language; even general programmers have less exposure to it and find it difficult to learn, not to mention computational chemistry programmers. Migrating the workflow to Rust generally means that C++, Fortran, Python or Matlab cannot satisfy us in some situation, and this dissatisfaction is not easily compromised:

• Matlab is commercial software and not a general-purpose language;
• Python has defects in performance and parallelism, and memory control is difficult; Python's MPI is not convenient;
• C++ is not impossible, but it is not easy to write code that correctly conforms to program standards, and we currently do not have senior architects; C++ project organization is not as convenient as Python's PyPI or Rust's Cargo; C++ templates and macros are too flexible, easily producing hidden compilation problems, and IDEs cannot correctly recognize some syntax.
• C is not a high-level language, and problems like compiling correctly but failing at runtime, segmentfaults, etc., easily occur;
• Fortran has similar problems to C in its low-level parts, and as a high-level language it is not as good as C++.

Rust indeed solves most of the above problems well; there are still some regrets in actual experience, but most regrets are not decisive.

One of the important reasons to choose Rust, I believe, is its higher performance. The most important reason for abandoning Python in favor of Rust is generally performance; the secondary reason is that Python's overly flexible program framework makes it easy for novice programmers to write code that does not conform to program standards, polluting the codebase. Therefore, in my view, for people who use Rust and abandon Python (for me, as someone who moved out of the comfort zone of the Python workflow to Rust), performance is non-negotiable, at least one cannot make great concessions in performance for higher development efficiency. Otherwise why not use Python.

I did try to write the Rust math library RSTSR. One of the important factors is the hope to advance the development efficiency of electronic structure programs in Rust. But this is advancing development efficiency based as much as possible on fast program performance, not the other way around.

6.2. How to Evaluate and Improve Computational Performance?​

Although Rust is said to have higher performance, it must be pointed out that

1. C/C++/Julia, as compiled languages, generally also have higher performance.
2. Language is not the only factor determining performance. Even with the same algorithm, the way the program is written (techniques) also determines program performance.

To avoid falling into agnosticism (making the discussion ineffective) when discussing program performance, it is necessary to state here that program performance has its evaluation strategy. This generally has nothing to do with Rust; as long as it is a language with a relatively powerful compilation backend (like GNU, LLVM), the discussion here applies.

Most of the understanding here took shape from implementing the CNN Winograd algorithm. The CNN Winograd algorithm is a very interesting example: it is both compute-intensive and requires heavy bandwidth usage; and to improve computational efficiency, assembly language (or instruction set functions close to assembly) is needed. The content discussed below is reflected in this document.

1. The premise of pursuing performance is an efficient algorithm. Algorithm has two meanings:

   • Algorithm complexity. A typical example is the MP2 4c-2e AO2MO problem:

     $$g_{ij}^{ab} = \sum_{\mu\nu\kappa\lambda} C_i^{\mu} C_a^{\nu} g_{\mu\nu}^{\kappa\lambda} C_{\kappa}^{j} C_{\lambda}^{b}$$

     If no optimization is done and the above formula is computed directly, an 8-fold loop would be written (because there are 8 indices $i j a b \mu\nu\kappa\lambda$), and each iteration does 4 multiplications and 1 addition, so the FLOPs of the above formula is $5 n_{\text{occ}}^2 n_{\text{vir}}^2 n_{\text{basis}}^4 \sim O(N^8)$ (i.e., 8th-power complexity). But in fact, without any approximation to the above formula, the same result with no error can be obtained with less computation:

     $$g_{ij}^{ab} = \sum_{\lambda} C_{\lambda}^{b} \sum_{\nu} C_{a}^{\nu} \sum_{\kappa} C_{\kappa}^{j} \sum_{\mu} C_{i}^{\mu} g_{\mu\nu}^{\kappa\lambda}$$

     This requires only 4 $O(N^5)$ operations (i.e., 5th-power complexity), greatly simplifying the computation time.

   • FLOPs, the floating point operation count. A simple example: for square matrices $\mathbf{A}, \mathbf{B}, \mathbf{C}$ of length $n$,

     $$\mathbf{A}\mathbf{C} + \mathbf{B}\mathbf{C}$$

     The FLOPs of this expression is $4 n^3$. But if we combine like terms to get $(\mathbf{A} + \mathbf{B})\mathbf{C}$, the computation can be halved to $2 n^3$. It does not reduce the complexity below cubic, but the performance improvement is still great. This looks like a very simple conclusion, but this method is of great help when dealing with the $O(n_{\text{occ}}^3 n_{\text{vir}}^3)$ terms of RI-CCSD. There is also a class of problems that can reduce computation using symmetry, such as RI-JK integrals, RI-MP2 energy summation, the pp-Ladder term computation of CCSD, etc.

   The above discussions do not involve concrete programs or concrete code techniques, but they are the most critical performance improvement factors. Before using code techniques to accelerate a program, one should as much as possible explore the possibility of improving the algorithm and the FLOPs.

2. Determine the FLOPs or bandwidth usage. Most electronic structure problems, especially those not exploiting sparsity and locality, can have their FLOPs given strictly. This does not mean computing exactly correctly, but at least estimating without deviating more than 10%.
   Taking the $O(n_{\text{occ}}^3 n_{\text{vir}}^3)$ terms of RI-CCSD as an example,

   $$\begin{aligned} I_{kldc}^{(1)} &= \sum_P B_{kd}^{P} B_{lc}^{P} \\ I_{ikca}^{(2)} &= \sum_{ld} I_{kldc}^{(1)} t_{il}^{da} \\ I_{ikca}^{(3)} &= \sum_{ld} I_{kldc}^{(1)} t_{il}^{ad} \\ I_{ikca}^{(4)} &= \sum_P \left(B_{ik}^{P} + \sum_c B_{dk}^{P} t_i^{d}\right) \left(B_{ac}^{P} - \sum_l B_{lc}^{P} t_l^{a}\right) \\ t_{ij}^{ab} &\leftarrow \sum_{kc} \left(I_{ikca}^{(2)} - I_{ikca}^{(3)} - I_{ikca}^{(4)}\right) t_{jk}^{cb} \\ t_{ij}^{ab} &\leftarrow \sum_{kc} \left(\frac{1}{2} I_{ikcb}^{(2)} - I_{ikcb}^{(4)}\right) t_{jk}^{ca} \end{aligned}$$

   The above has 6 lines of formulas, but the computation of $I_{kldc}^{(1)}, I_{ikca}^{(4)}$ is actually negligible; the remaining four lines are each $2 n_{\text{occ}}^3 n_{\text{vir}}^3$ FLOPs, so the total computation of this process is estimated to be $8 n_{\text{occ}}^3 n_{\text{vir}}^3$. Is the rest unimportant? Not very. Looking carefully, $I_{kldc}^{(1)}$ can exploit symmetry, so the FLOPs of $I_{kldc}^{(1)}, I_{ikca}^{(4)}$ are about $3 n_{\text{occ}}^2 n_{\text{vir}}^2 n_{\text{aux}}$; its ratio to the computation of the other 4 lines is about $1 / n_{\text{vir}}$, i.e., when the number of virtual orbitals is greater than 10, the computation of $I_{kldc}^{(1)}, I_{ikca}^{(4)}$ can be ignored.
   The above formula also has some memory bandwidth usage, reflected in the computation of $\left(I_{ikca}^{(2)} - I_{ikca}^{(3)} - I_{ikca}^{(4)}\right)$. Note that for a computing device with one NUMA node and 16 cores CPU, the typical bandwidth and performance limits are 10 GB/sec (1.2 G doubles/sec) and 1 TFLOP/sec respectively. If the FLOPs and the memory bandwidth demand differ by more than 800 times, then memory bandwidth is not the decisive factor. In the current problem, if $n_{\text{occ}} n_{\text{vir}} > 800$, then memory bandwidth no longer matters either. Nevertheless, memory bandwidth analysis is still important in some specific memory-intensive problems (like DFT computation), or in MPI communication.

3. Determine the performance limit of the computing device. Program performance optimization always has an upper bound; this upper bound is not unknowable either. By running the bandwidth64 program in parallel, one can generally evaluate the computation and communication efficiency of the L1, L2, L3 caches of the computing device. Intel OneAPI Advisor can also provide valuable data. Generally speaking, current CPUs usually have 70 GFLOP/sec/core (depending on the CPU frequency and the number of AVX channels), and the bandwidth is 16 GB/sec/socket (which should depend on the motherboard and the CPU PCIe). For this part of the information, one can also refer to the previous evaluation document of the HiSilicon Kunpeng device performance. GPU computational performance parameters are generally more transparent than CPU, and can be looked up directly in the official manuals.

4. Compare the device limit and the actual program's floating point efficiency FLOP/sec. Taking the RI-CCSD Rust demonstration case as an example, we can determine that the most time-consuming part averages about 45% utilization of the device performance limit (around 500 GFLOP/sec on a 16-core CPU device). Generally speaking, a program efficiency of 50% of the device performance limit is satisfactory:

   • On the one hand, high-performance BLAS achieves only 60%–80% utilization, and further optimization hardly leaves room for improvement;
   • On the other hand, computational chemistry currently has low output value and is usually exploratory research; invalid computation and inefficient usage scenarios are common, so it is acceptable as long as the efficiency is not unbearably low (about 50% extra efficiency loss, i.e., 20% of the device performance limit);
   • Moreover, part of electronic structure tasks are iterative; more iteration steps are themselves performance loss, so pursuing extreme performance also includes pursuing better numerical iteration algorithms.

   For Python-based programs, considering only the single-node situation, the efficiency is often tolerable in many cases; but there are also many cases where Python and NumPy step on the intolerable low-efficiency red line (this is somewhat why PySCF has so much C code patching).

6.3. As Chemistry Programmers, Do We Need to Manually Tune Performance?​

Generally speaking, we should not. Chemists, including program engineers working in chemistry, should not waste energy on the showmanship of program optimization, but should let the program return to its essence as a "formula translator". The premise of this statement is that the algorithm has already been designed by chemists and engineers.

Everyone's energy is limited; if everyone can handle their own work well, efficiency will improve quickly. We chemists had better hand over the "formula translator" work entirely to professional numerical program architects, and focus on our own work. But this also means that the "formula translator" itself must be good enough, and professional numerical program architects must be able to do their own work well, for us to use it conveniently.

As a more detailed explanation, the main FLOPs consumption of electronic structure programs is in matrix multiplication; and matrix multiplication is generally provided by BLAS, which we have no ability to manually tune. For the rest, slightly worse performance is acceptable, but like Python and NumPy it is still not very satisfactory. So the room for us to manually tune performance is actually not large; for the three-level caches, pipeline streaming, and prefetching that high-performance computing cares most about, we treat them all as if they did not exist; we write some parallelism when appropriate. This approach is generally fine on CPU, but GPU may need further discussion.

But specifically, this question depends on how "manual performance tuning" is defined, and under what toolchain the tuning is done.

• One extreme case is that I say I like np.einsum (or equivalent tensor contraction tools in other libraries), and I do not allow more complex expressions. But considering that current tensor contraction tools cannot reach limit efficiency (though it must be admitted that many tensor contraction tools can now achieve ideal efficiency, at least higher than 20% of the limit performance); and some expressions written as Einstein summations may not be suitable for combining like terms, which indeed increases computation time; then manually splitting Einstein summations and combining like terms is necessary.

• A more common but hidden case is that we may encounter non-contiguous dimensions. For example the following problem:

  $$\mathcal{T}_{jbP} = \sum_{ia} B_{iaP} I_{ijab}$$

  Although this problem can indeed be implemented by splitting into GEMM (so Einstein summation may also have considerable efficiency):

  # T: jbp, B: iaP, I: ijab
  T = np.zeros([nocc, nvir, naux])
  for j in range(nocc):  # can be parallel
      for i in range(nocc):  # can be parallel, but with reduce
          T[j] = I[i, j].T @ B[i].T
  

  But if $I_{ijab}$ is not stored in the order $i, j, a, b$, but rather as $I_{iajb}$, then not only is the program more efficient, it can also be done with shorter and clearer code logic:

  $$\mathcal{T}_{jbP} = \sum_{ia} B_{iaP} I_{iajb} \quad \Rightarrow \quad \mathcal{T}_{jb, P} = \sum_{ia} B_{ia, P} I_{ia, jb}$$

  # T: jbp, B: iaP, I: iajb
  T = I.reshape([-1, nocc * nvir]).T  @ B.reshape([-1, naux])
  T.shape = [nocc, nvir, naux]
  

  Therefore, before generating the four-dimensional tensor $\mathbf{I}$, we should already consider whether it is more suitable to store it as $I_{ijab}$ or $I_{iajb}$. More extreme problems may not even be expressible with GEMM in any way (an explicit transpose must be performed to convert to a GEMM problem), such as $\mathcal{T}_{Pca} = \sum_b B_{abP} f_{bc}$ — such tensor storage patterns must be avoided.
  This type of problem must be considered when writing programs, but it is not something chemists themselves should care about. The "formula translator" itself is only responsible for translating formulas correctly, not for what is the most efficient arrangement of tensor indices. This kind of problem is in the gray zone between chemists and the "formula translator", and needs to be solved by the engineers implementing electronic structure methods.

• Whether the toolchain is mature, efficient, and whether there are problems specific to computational chemistry, are also factors to consider. Math library quality is not high; although this is nominally the responsibility of numerical program architects, chemists may need to lend a hand because there is no one else. Take NumPy as an example here.

  • Apart from fast matrix multiplication, NumPy's other operations are inefficient; this caused PySCF to spend considerable effort on DFT-related grid integration problems in C. This is partly NumPy's own problem (it cannot parallelize vector operations well). Early NumPy's matrix multiplication probably had problems too; PySCF's np_helper.c has had to clean up after NumPy quite a lot.
  • NumPy does not provide a GEMM interface itself; for complex matrix $\mathbf{C} = \mathbf{A}^{\dagger} \mathbf{B}$ multiplication, one has to call SciPy to solve it (this problem is also not easy to solve in the current RSTSR framework in Rust). So whether to call a.conj() @ b or scipy.linalg.blas.gemm, and whether the matrix fed to SciPy has been transposed to f-contiguous, these fine-grained tricks indeed affect program efficiency to some extent.
  • We often store symmetric matrices as lower triangular matrices (c-contiguous) or upper triangular matrices (f-contiguous); but when actually using them, they need to be expanded to symmetric matrices. This feature is frequently needed in chemistry, and many mainstream math libraries do not support this type of feature well. Expanding such packed matrices is fine to implement in compiled languages, but implementing it in Python is inappropriate; therefore chemists must develop program tools to solve this kind of problem.

6.4. I Cannot Trust Math Library Performance for Some Problems; Should I Write My Own Operators?​

My answer is that at present it is likely still necessary; but this depends on the completeness of the math library.

I do not deny that some math libraries have very bad efficiency on some simple operation problems. Not being able to trust math library performance is completely understandable. I cannot give a firm answer to this question.

But I also think that on many simple standard problems, it is reasonable for math libraries to have good performance. This should not be left to users. Bad math library performance is the math library's problem; as users, we only need to clean up their mess because there is no way around it.

There are also many non-standard problems, especially in machine learning where operator fusion was a hot topic for a while; such problems can usually be reduced to standard problems, but with the necessity of performance optimization or numerical stability optimization (a typical example is softmax). Such problems should have first been solved manually by machine learning deployment engineers; but later PyTorch solved the problems by throwing manpower at operators.

In electronic structure problems, especially the CC/MP$n$ class of problems, everything is matrix multiplication, elementwise multiplication or addition, and summation problems. Such problems are very standard, and there are rarely special computation patterns. Even non-standard problems are usually not FLOPs bottlenecks.

Matrix multiplication should trust BLAS libraries. Even using MKL on AMD CPUs, we can absolutely never write anything faster than BLAS. As a compute-intensive algorithm that is unfriendly to cache access, the concrete algorithm of matrix multiplication is not only complex, but requires very subtle tricks. Besides the complex five-fold loop exploiting the three-level caches (the batch of the loop depends on the L1/L2 cache of the current computing device), it also needs assembly language combined with pipeline streaming and prefetching for the corresponding CPU microarchitecture; this is no longer a problem we should understand. As an introduction (I have not fully understood it either), please refer to the tutorial of R. van de Geijn, the postdoc advisor of Kazushige Goto (K. Goto) and Devin Matthews.

But many DFT problems, although they can also be implemented with efficient math library functions, may be better without the math library. DFT's performance bottleneck is different from those of HF and post-HF; although two of its steps are GEMM matrix multiplication problems:

and the basis transformation problem of orbitals on grid points:

Apart from these two GEMM problems, the other problems are usually not GEMM problems; they are usually bandwidth bottlenecks, but also occupy considerable computation time. Among such problems, some can be efficiently implemented in NumPy with np.einsum, but np.einsum cannot guarantee high efficiency on all problems.

6.4.1. An Operator Case Not Suitable for Trusting the Math Library: $C_g = \sum_i A_{ig} B_{ig}$​

A typical problem is the generation of the SCF density on grid points:

One step of this problem can be reduced to the following numerical problem (based on some algorithms, matrix $A$ can be equal to matrix $B$, further saving bandwidth):

This step can be implemented in a standard way in a math library:

Written as NumPy code:

t = a * b
c = t.sum(axis=0)

Of course, the above code can also be abbreviated as c = (a * b).sum(axis=0). Such an algorithm is of course correct, but the cost is generating a temporary tensor $T_{ig}$. The extra memory needed for this temporary tensor is a small matter (because DFT should generally control the number of batched grid points); but the extra write and read of $T_{ig}$ is a waste of computing resources.

Next, we will use a concrete example to show and compare the computational efficiency of different math libraries and different implementations. For the current problem, the number of orbitals (index $i$) is set to 1000, and the number of grid points (index $g$) is set to 100,000.

We first consider the optimal solution of this code. In Rust, directly computing with a brute-force loop:

let vec_c: Vec<f64> = vec![0.0; ng];
vec_c.par_iter_mut().enumerate().for_each(|(g, c)| {
    for i in 0..ni {
        *c += vec_a[i * ng + g] * vec_b[i * ng + g];
    }
});

This computation takes 31 msec. Note that matrices $\mathbf{A}, \mathbf{B}$ each have a memory size of 0.75 GB; traversing these two matrices alone takes about 15 msec; therefore the above implementation reaches at least 50% of the limit bandwidth performance. I cannot yet determine whether there is a faster implementation for the above computation. PySCF has a slightly different algorithm for this problem (the C function VXC_dcontract_rho), which takes 29 msec. I think taking 30 msec as a reference is reasonable.

Do not think that because the fastest for this problem is only 30 msec, it does not need attention. On the one hand, there are quite a few similar processes in DFT computation; on the other hand, if this function is not implemented well, it can take seconds, as will be seen below.

NumPy actually has another tool for this kind of problem: np.einsum:

c = np.einsum("ig, ig -> g", a, b)

This function actually has very good performance on the current problem, taking about 36 msec. Moreover, this function runs completely single-threaded, meaning it likely has very strong SIMD optimization but is not parallelized. Due to the characteristics of the current problem, not parallelizing can still achieve good efficiency; but not every problem is like this, and np.einsum will show its disadvantages on other problems.

If we cannot use np.einsum, but use ordinary tensor library logic, then for NumPy:

# numpy                       #  100 msec in total
t = a * b                     #   79 msec
c = t.sum(axis=0)             #   21 msec

For Rust's RSTSR:

// RSTSR                      // 127 msec in total
let t = &a * &b;              //  97 msec
let c = t.sum(0);             //  31 msec

For Rust's ndarray:

// ndarray                    // 340 msec in total
let t = &a * &b;              // 308 msec
let c = t.sum_axis(Axis(0));  //  21 msec

For Rust's nalgebra:

// nalgebra                   // 410 msec in total
let t = a.component_mul(&b);  // 343 msec
let c = t.row_sum();          //  70 msec

Although NumPy and RSTSR have considerable efficiency, they are in any case 3–4 times slower than the most efficient code. Therefore, for the current $C_g = \sum_i A_{ig} B_{ig}$ problem, to achieve efficiency, it is very likely necessary to implement it in the Einstein summation framework or to hand-write the function, and it cannot be implemented with ordinary matrix multiplication, vector elementwise operations, summation and other basic operators.

6.5. Which Rust Features Are Unsuitable for Building Math Libraries?​

• Rust is not a language focused on scientific computing. Things like transpose, inversion, etc., can neither be replaced by symbols (operators) nor implemented in the language standard library; even complex number operations require external libraries.

• Rust is difficult to implement syntactic sugar with symbols. Some operator-corresponding traits in Rust are hard-coded; this causes functionality that can be concisely implemented with operators in NumPy to have to be implemented with functions in Rust, greatly increasing syntactic noise. For example,

  • Indexing. In NumPy, a[1, 2:4, 3:10] can index a three-dimensional tensor to a two-dimensional matrix; this indexing is done through the operator []. Rust also has the operator [], but it is done through the Trait Index:

    pub trait Index<Idx>
    where
        Idx: ?Sized,
    {
        type Output: ?Sized;
    
        // Required method
        fn index(&self, index: Idx) -> &Self::Output;
    }
    

    Its output is &Self::Output, not Self::Output! This means you cannot only return an existing variable, but cannot return some created quantities. We will discuss later that the raw data of an indexed tensor is indeed unchanged, but new shape information and offset are produced; these need new memory space to store. Therefore, the return value of tensor indexing is in any case very difficult to store with a reference type &Self::Output. Thus, for tensor indexing in Rust, unless you just want to take out a value, which can be done with the Trait Index or the equivalent operator []; otherwise, if you want to index out a sub-tensor, you still have to honestly use a function. In the candle framework this function is .i(), in ndarray it is .slice(); in RSTSR, .i() and .slice() are equivalent.

  • Assignment. C++'s symbol overloading is so powerful that C++ tensor contraction libraries can do wonders with the operator =. But Rust is the other extreme: = only has assignment functionality and cannot be overloaded. This may be designed for Rust's lifetime guarantees, but it is very unfriendly to us scientific computing users. The simplest example is assigning to an indexed tensor:

    c[:, :nocc] = a[:, :nocc]
    

    Such code is very intuitive in Python, but Rust does not allow assignment to a statement. In Rust's RSTSR, although this is not unsolvable, it still goes around in circles or adds syntactic noise:

    // by assign
    c.i_mut((.., ..nocc)).assign(a.i((.., ..nocc)));
    // if c is zeroed before assignment, use add_assign is also valid
    let mut c: Tensor<f64> = rt::zeros(([nao, nmo], &device));
    *&mut c.i_mut((.., ..nocc)) += a.i((.., ..nocc));
    

    Please note that the Trait AddAssign in Rust is overloadable, so the above problem has two solutions. I personally prefer the second one; since it can use the += symbol, it is more like Python code, but it has a performance cost. The same is true for ndarray.

  • Comparison. Taking the function eq as an example, in the standard library's PartialEq,

    fn eq(&self, other: &Rhs) -> bool;
    

    Its return type is hard-coded to boolean. But in NumPy, code like this

    c = a == b  # returns a tensor of boolean type, not bool itself
    c = a is b  # this returns boolean, but comparing id(a) and id(b)
    

    especially the first line above, must in Rust be implemented through other functions rather than symbols.

• Rust adds syntactic noise due to its ownership mechanism and error handling mechanism. This problem is also discussed in the standard library; refer to the Vec::try_with_capacity(_in) issue. This is an indirect rather than direct consequence. Rust's ownership and error handling mechanisms are actually its advantages; but as a library developer, when you start writing functions, you will find that you need to write at least 4 cases. Taking transpose as an example:

  • the function into_transpose_f taking values in and out and allowing error handling;
  • the function into_transpose taking values in and out and panicking directly on error;
  • the function transpose_f taking references in, outputting TensorView, and allowing error handling;
  • the function transpose taking references in, outputting TensorView, and panicking directly on error;

  As a user, generally only transpose is used, same as NumPy. But as a library developer, writing two functions for the two cases of taking references or taking values is something that should be done; you do not know whether users have special requirements for error handling, so it is best to handle the error-handling case (in naming, RSTSR and tch-rs use the _f suffix, and many standard library functions use the try_ prefix). And reshape or to_layout functions need to handle two more cases: outputting a Copy on Write type or outputting a value type. So although Rust is very powerful, library developers have to make at least four functions every time they write a feature, many of which are not easy to use with macros; library users also at least need to know which of transpose and into_transpose is the value type and which is the reference type, which is quite laborious.

• Overloading with Traits is not very convenient. The Rust language itself forbids override and overload. As a non-OOP language, forbidding override is absolute; but as a generic language, overload can be implemented in a roundabout way with generics. Taking the asarray function as an example, we can first define the trait

  pub trait AsArrayAPI: Sized {
      type Out;
      fn asarray(self) -> Self::Out;
  }
  

  then specialize for various tuple types. For example, if we want to get a one-dimensional tensor from a Vec<T> input, then we specialize for the Vec<T> type (as a 1-element tuple):

  impl<T> AsArrayAPI for Vec<T> {...}
  // this allow usage of `asarray(vec)`
  

  But if we want to directly get a two-dimensional matrix from the input tensor, then we specialize for the 2-element tuple (Vec<T>, Layout<D>):

  impl<T, D> AsArrayAPI for (Vec<T>, Layout<D>) {...}
  // this allow usage of `asarray((vec, layout))`
  

  As a concrete use case,

  use rstsr_core::prelude::*;
  use rstsr_openblas::DeviceOpenBLAS;
  let vec_a = vec![0.0; 15];
  let vec_b = vec![0.0, 15];
  let vec_c = vec![0.0, 15];
  let device_openblas = DeviceOpenBLAS::default();
  // generates 1-D tensor of shape [15] on default device (DeviceFaer)
  let tensor_a = rt::zeros(vec_a);
  // generates 2-D tensor of shape [3, 5] on default device (DeviceFaer)
  let tensor_b = rt::zeros((vec_b, [3, 5]));
  // generates 2-D tensor of shape [3, 5] on DeviceOpenBLAS
  let tensor_c = rt::zeros((vec_c, [3, 5], &device_openblas));
  

  This does achieve the goal of overloading with the trait system, and to some extent it is even more flexible than Python (Python does not allow overload on the surface, but allows optional parameters), but there are at least two problems:

  • Tuple type parameters must be input, i.e., implementing these functions requires at least two brackets. But for 1-element tuples, the two brackets are redundant, and the VSCode editor will remind you about it. This adds syntactic noise and creates inconsistency in code style; but this is much better than defining three functions asarray_to_1d, asarray_with_shape, asarray_with_shape_and_device. In Python, indexing has the syntactic sugar tensor[(a, b, c)] equivalent to tensor[a, b, c]; I wonder whether Rust has such a plan.
  • Defining an overridable function becomes very complex. It was originally a matter of declaring the function signature a few more times, but now an extra trait is needed. This is an additional burden on API developers. And I have no idea how the API documentation should be written either.

• The friendliness of API documentation is questionable. This is a rather subjective judgment. API documentation is not necessarily well done in any language or any framework. We can see quite good documentation for many Python or C++ libraries, but there is also a lot of manual effort in them. Cargo doc gives us little room to intervene; it guarantees the lower bound of API documentation, but also constrains its upper bound, and may not be suitable for large projects.

• The trait system still has room for improvement. Rust's trait system is easier to write in a standardized way than C++, solving problems before compilation; this is the benefit of Rust's trait system. But the downside is that there are many restrictions. As an example, when implementing the abs absolute value function, we notice that num::Signed and num::complex::ComplexFloat have two abs implementations. As a math library, we need the abs function for integer types (only implementing Signed), for floating point types (implementing both Signed and ComplexFloat), for complex floating point types (only implementing ComplexFloat), and perhaps also for unsigned integer types (only implementing Unsigned). Although Rust allows you to do this:

  impl<T> DeviceAbsAPI<T> for DeviceCpuSerial
  where T: Signed {...}
  

  But if you also want to implement ComplexFloat at the same time:

  impl<T> DeviceAbsAPI<T> for DeviceCpuSerial
  where T: ComplexFloat {...}
  

  Rust will tell you there is a conflict implementation. First, because floating point types implement both Signed and ComplexFloat, there is indeed an implementation conflict. Second, you will find that floating point numbers cannot naturally be Unsigned, so I should be able to implement for Unsigned, right? But this is also not allowed, because Rust does not know which downstream user, on which day, for what ulterior purpose, created a type myf64 and implemented both Unsigned and Signed for the myf64 type, so that your library really has a conflict implementation. Therefore, there are now three strategies for implementing the abs absolute value function:

  • directly specialize for concrete types, without using where clauses;
  • declare two traits (DeviceRealAbsAPI, DeviceComplexAbsAPI), both implementing the abs function;
  • create a new trait of your own, unifying the usage of the abs function.

  RSTSR currently adopts the 2nd strategy; but I can hardly say this is a good approach, because it either splits floating point types from complex floating point types, or splits integer types from floating point types. I believe this is also the perplexing problem of the crate num (discussed in issue #64).
  The 3rd strategy is actually not impossible, but it has another problem: should we optimize memory for by-value unary operations? If the type is f64, whose absolute value is also f64, then the following operation is feasible:

  a = a.abs()
  

  This does not allocate new memory to store variable a. But if the type is Complex<f64>, after taking the absolute value it is no longer the original type but f64, meaning with normal code style, the original memory cannot be reused. Therefore, if we want to reuse memory as much as possible, even if f64 and Complex<f64> can be unified with the same trait interface, they must take different code paths.
  I do not know whether such problems can be implemented in the future with negative bounds (so that there can be trait restriction strategies based on the inclusion-exclusion principle, instead of the current overly strict orphan rule).

6.6. My Misgivings about Iterators​

I have doubts about whether iterators are more efficient than for loops.

For many problems, writing with for loops is quite intuitive, and the efficiency loss is not large, if there really is any efficiency loss. Moreover, if iterators were really faster, I believe rustc would recognize the for loop and package its loop part into an FnMut closure, compiling it in the iterator way.

In my view, iterators do have meaning in the following uses:

• Iteration of non-1-D vectors. A tensor structure is not a one-dimensional vector, and there are many different possible iterations; and if written with for loops, it becomes multiple loops (with the number of layers not known at compile time). In this case, iterators are a rigid requirement.
• Parallelism. When a programming language cannot support parallelism at the syntactic sugar level (like Julia etc.), or cannot support it through macros or precompilation directives (like OpenMP etc.), then functions must be used. Functions must act on objects, so the loop itself must be packaged as an object; this object is the iterator. Rayon iteration, I believe, also came about this way by our scholarly predecessors.

There are times when iterators may be more convenient to use, such as the collect::<Vec<T>>() scenario; otherwise we have to create an empty mut Vec<T> outside the loop and add a push(item) statement in the loop body, which is rather troublesome. But this is limited to cases where the loop logic is very simple; if the loop logic is somewhat more complex, for loops are more intuitive instead; and besides, everyone learned C/C++/Python first, and is more used to for loops.

Therefore, when iterators are not a rigid requirement, I do oppose or support using iterators out of code habits (rather than more rational reasons); especially without micro benchmarks testing and proving whether iterators are more efficient than for loops.

At the same time, it should be pointed out that for the problems encountered in electronic structure, iterators can generally be designed to be efficiently parallel; but this costs development time and the cost of mutual adaptation among programmers, and program efficiency may not really lack that little bit, since the main time consumption comes from BLAS; apart from that, the key to improving program efficiency is still the design of the algorithm itself. Therefore, to be honest, I hold a reserved attitude toward designing special iterators; I tend to use for loops; on this point, if I start to deeply understand the bottlenecks of self-consistent field algorithms in the future, I will revisit this problem.

Two iterator-related problems are discussed here.

• Parallel iterator design. This is the content of the Rayon library, not my experience. First, my experience is that Rayon's parallelism is in some sense more powerful than C OpenMP; it at least supports double iteration loops, has flexible function interfaces similar to MPI, and in terms of performance, at least in simple problems, I do not see an obvious difference from OpenMP. For OpenMP's simple for loops, Rayon can directly convert to the equivalent

  (0..n_i).into_par_iter().for_each(|i| {...})
  

  Although there is some syntactic noise, it is indeed a one-line solution. Unlike OpenMP, the parallel loop body must be Fn rather than FnMut, and cannot pass mut or types that do not implement the Send trait; it may require unsafe or other support to implement the complete OpenMP-like shared parallel mode.

  We sometimes need to iterate some structures in parallel, and such structures cannot be solved with simple for loops (such as iterating over all tensor elements). In this case, although implementing Iterator can also be parallel, the bridge mechanism is needed; bridge is very inefficient. To make iterators more efficient, Rayon requires us to provide the following trait implementations:

  • DoubleEndedIterator, i.e., the double-ended iterator;
  • ExactSizeIterator, i.e., the iterator length needs to be known in advance;
  • Producer, which needs to implement the split_at function. An iterator can be seen as a mapping, i.e., a one-to-one correspondence between the iteration list and the continuous integer array. Rayon parallelism needs the mapping from the integer array to the iteration list (but not its inverse); i.e., from the $x$-th iterated element, its iterated content can be deduced. Taking a two-dimensional matrix of size $(i, j)$ as an example, the element corresponding to the $x$-th index in row-major is $(\lfloor x / j \rfloor, x \bmod j)$.
  • ParallelIterator and IndexedParallelIterator; these are the high-level traits in Rayon.

• Use array [usize; N] as much as possible and avoid slice [usize]. This is a problem encountered in high-dimensional tensor indexing, not a problem common to all computational chemistry or programs. We will encounter this problem in $C_{ij} = A_{ij} + B_{ji}$, i.e., the matrix transpose summation problem.

7. RSTSR Features​

7.1. What Are RSTSR's Design Goals?​

• Function style: as consistent with NumPy as possible, as simple and easy to use as possible and conforming to NumPy habits.
• Functionality: most of the NumPy main program, SciPy's scipy.linalg and scipy.linalg.blas.
  • Most math libraries in other compiled languages are also benchmarked against NumPy, represented by xTensor (C++), ndarray (Rust), GoNum (Go) and NumSharp (C#). NumPy is of course not the best math library, but there is nothing shameful in using it as a reference.
• CPU performance: matrix multiplication connects to BLAS; most of the other operators are faster than serial NumPy under 8-core parallelism. It allows flexible integration of Rayon parallel code.
• Backends: hope to support GPU programming in the future (including CUDA, HIP).
• Types: support data storage of all types through trait generics; support computation of some types through std::ops and the crate num.

Electronic structure programs not only pursue the above goals, but also need to implement electronic structure algorithms themselves, and even use algorithms not available in standard math libraries. My knowledge breadth is not enough, but I previously heard that a student of Prof. Ren Xingguo (poster) used oneMKL PARDISO when implementing basis-set-limit RPA. Obviously the features chemistry needs are broader than those of general math libraries. Therefore, even just for writing electronic structure programs, these simple mathematical program tools are worth developing.

For RSTSR's goals and the survey briefing on other Rust or non-Rust math libraries, the RSTSR 1st report has more detailed descriptions.

At the same time, it should be noted that although RSTSR is designed for electronic structure program purposes, it is still a math library, not an electronic structure library. Electronic structure programs can separate part of their mathematical problems into a self-contained system, and the problems that can be solved are not limited to electronic structure; therefore we should not force the claim that RSTSR is a math library with electronic structure computing features. Our vision can and should be longer-term.

7.2. What Features Does RSTSR Currently Have? What Is Still Lacking for Electronic Structure?​

With a total of about 3 months of development time (7/28–8/5, 8/11–10/6, 12/22–1/21, during which other work such as the RHF first-order gradient implementation draft under REST was completed), RSTSR has currently completed most of the basic features required of math libraries by the Python array API standard. Please refer to the electronic structure needs for math libraries in What math library features do electronic structure programs need?.

I myself also, based on the following three-step workflow: 1) drafting with NumPy np.einsum, 2) removing np.einsum in NumPy to prepare for migration to Rust, 3) concretely implementing the algorithm based on the math library RSTSR, starting from the non-RI closed-shell CCSD tutorial document, and essentially without drawing on derivations in existing literature, implemented closed-shell RI-CCSD in Rust (showcase_rust_riccsd), taking less than 4 days in total (2025/1/17–2025/1/20), with program performance at least comparable to the Psi4 fnocc module, though the memory overhead was not yet more carefully optimized. I hope this project can show that the NumPy + Rust workflow strategy, aided by a relatively complete Rust math library, is entirely capable of achieving efficient and rapid implementation of electronic structure methods.

RSTSR still has many deficiencies to solve.

• It does not satisfy all function requirements of the Python array API standard, including concat, stack, argmax, etc.
• The BLAS and Lapack wrappers and the linear algebra part are not implemented; this means RSTSR currently has no definite solution for how to solve matrix eigenvalue or vector solving problems. Therefore, the current RSTSR cannot yet do SCF, CP-HF, Cholesky decomposition, or the matrix logarithm or determinant computations needed by RPA.
• There are no correctness and efficiency tests. Before starting the RI-CCSD implementation, I kept adding features to RSTSR and had not yet done any correctness verification. So it is entirely possible that RSTSR has numerical errors. Of course, during the RI-CCSD implementation, there was only one rather benign bug fix (removing an overly strict trait bound), so I still hold great expectations for program correctness.
• Some advanced indexing features of NumPy are not implemented, especially indexing tensors by integer lists.
• Heterogeneous backends are not implemented. Although multiple backends are currently implemented, they are generally only used to switch between different BLAS engines, verifying the possibility of multi-backend implementation. A true multi-backend should consider CUDA and HIP, and for MacOS, Accelerate and Metal.
• More implementation and testing on more electronic structure problems is still needed.
• Adaptability under MPI parallelism has not been verified.

7.3. Why Not Be Satisfied with Existing Tools? Which RSTSR Features Are Advantageous?​

• Supports complete n-dimensional arrays and their broadcasting. This is relative to ndarray; it only supports partial broadcasting and manipulation functions, and should be considered half-finished. Many other Rust math libraries do not support n-dimensional arrays, notably nalgebra, faer, etc. There are also libraries that do not support dynamic-dimension tensors (possibly because of computational graph construction), notably dfdx, burn, etc.

• Supports floating point types and complex floating point types. Rust's current machine learning libraries generally do not support complex floating point, notably candle, burn, etc. To be honest, candle is the framework closest to a tensor library usable for computational chemistry, but its design pattern is the same as PyTorch's, not very extensible (although it was an excellent design 10 years ago), with types and backends hard-coded in the core program. This conflicts considerably with my development philosophy. In extensibility, we are closer to burn; but burn does not satisfy our needs in data types and variable dimensions.

• We plan to support multiple backends. This is relative to NumPy and ndarray. nalgebra abandoned its CUDA backend a few years ago. PyTorch is now invincible; if our own functions cannot beat PyTorch at anything, perhaps we can make PyTorch our backend in the future. But no matter which engine is the backend, basic arbitrary-dimension tensors, reshape, transpose and other features still need to be implemented in Rust; C++ cannot help with this.

• CPU parallelism of internal operators. This feature means that apart from BLAS computations that can be accelerated with ready-made libraries, the performance of the rest is also considerable; but we have not spent much energy on this, and obviously the limit efficiency is not yet reachable, but at least we will not be dragged down too much. Almost all math libraries represented by NumPy and ndarray only parallelize BLAS, and do not parallelize elsewhere. But a small number of math libraries with this feature have better performance than ours, represented by PyTorch.

• Seamless external Rayon parallelism. Most functions of this library can run inside Rayon threads.

• BLAS supports serial and parallel calls. This is learned from rest_tensors. Inside Rayon parallelism, the number of BLAS cores will be limited to one; without a parallel environment, it runs at full CPU; the number of concurrent cores can also be controlled in the device. In my understanding, ndarray and nalgebra do not support this feature; this type of feature requires specialization for each BLAS distribution to implement, rather than relying on the crate blas-sys as a one-size-fits-all solution; otherwise problems easily occur when calling BLAS in parallel with Rayon.

• A concise interface form close to NumPy. Taking the sum function as an example, our sum function requires passing the list of dimensions to be summed; for example, when computing the dipole moment, we need to compute the elementwise multiplication and summation of the one-electron integrals and the density matrix:

  // tsr_int1e_r: (t, mu, nu); rdm1: (mu, nu)
  let dipole = (tsr_int1e_r * rdm1).sum((-1, -2))
  

  But in ndarray, it needs to be written as

  let dip = (tsr_int1e_r * rdm1).sum_axis(Axis(2)).sum_axis(Axis(1));
  

  Not only is the code more complex, the computation cost also increases. Another example is the asarray function, which we overload using the trait overload pattern; but ndarray requires using many functions like from_vec, from_shape, from_shape_vec, from_shape_ptr to do one thing, which is unfriendly to users.

7.3.1. An Operator Where RSTSR Has Some Advantage: Elementwise Multiplication with Transposed Matrix $\mathbf{C} = \mathbf{A} \odot \mathbf{B}^{T}$​

Previously, on the operator $C_g = \sum_i A_{ig} B_{ig}$ problem, we were thrashed by NumPy's np.einsum; now it is time to win one back.

In MP2 computation, the following computational problem appears:

This operator is actually a fused operator of elementwise multiplication followed by summation, and it is best to implement it by the original formula; but for implementation convenience, for each pair of occupied orbitals $(i, j)$, two matrix elementwise multiplications are done: one without transpose, one with transpose:

This is the origin of the operators $\mathbf{C} = \mathbf{A} \odot \mathbf{B}$ and $\mathbf{C} = \mathbf{A} \odot \mathbf{B}^{T}$ we are examining.

Logically this operator should not take much time, since it is quadratic in time (multiplied by the pairs of occupied orbitals, the FLOPs is $O(n_{\text{occ}}^2 n_{\text{vir}}^2)$, much smaller than the 5th-power complexity $n_{\text{occ}}^2 n_{\text{vir}}^2 n_{\text{aux}}$ of the whole RI-MP2).

But if implemented with NumPy, the efficiency becomes very bad. First, RI-MP2 is naturally easy to parallelize over pairs of occupied orbitals $(i, j)$; but constrained by Python syntax, in pure Python we should not think about parallelizing for loops. NumPy's matrix multiplication uses BLAS, so even without parallelizing over pairs of occupied orbitals $(i, j)$, obtaining the MO-basis 4c-2e ERIs from Cholesky-decomposed ERIs is still fast; that is, NumPy's efficiency in the computational bottleneck part is guaranteed.

But for computing the elementwise matrix multiplication of the non-bottleneck part, NumPy has efficiency problems. Now suppose a system with 512 electrons, $n_{\text{occ}} = 256$, $n_{\text{vir}} = 1024$, $n_{\text{aux}} = 3072$; on my 16-core CPU laptop, assuming BLAS can reach 55% performance (600 GFLOP/sec), then the time for the computational bottleneck part of RI-MP2 is 328 sec.

So it seems the energy summation computation should not be the bottleneck, right? But running the code, you will find that for NumPy, the $1024 \times 1024$ matrix computations

• normal elementwise multiplication $\mathbf{C} = \mathbf{A} \odot \mathbf{B}$ is 250 µsec,
• transposed elementwise multiplication $\mathbf{C} = \mathbf{A} \odot \mathbf{B}^{T}$ is 5.3 msec,
• matrix summation $e = \text{sum}(\mathbf{C})$ is 130 µsec, needing to be computed twice,

combined with the loop over $n_{\text{occ}}^2 / 2 = 32768$ pairs of occupied orbitals, the above computation is estimated to take 190 sec. This already reaches 60% of the bottleneck computation time, not a negligible amount.

In NumPy, the part with efficiency problems is the transposed elementwise multiplication $\mathbf{C} = \mathbf{A} \odot \mathbf{B}^{T}$, and the summation $e = \text{sum}(\mathbf{C})$ is also inefficient; but the efficiency that can be gained from summation is limited. In Rust, we can naturally parallelize this problem over pairs of occupied orbitals $(i, j)$, thus alleviating the problem. This is one solution.

But even without parallelizing over pairs of occupied orbitals $(i, j)$, with only ordinary for loops, RSTSR can achieve very good performance on the RI-MP2 energy summation.

In RSTSR, if the input $\mathbf{A}$ and $\mathbf{B}$ are of fixed dimension Ix2, executed in parallel,

• the transposed elementwise multiplication $\mathbf{C} = \mathbf{A} \odot \mathbf{B}^{T}$ can be improved to 310 µsec (ndarray is 3.7 msec),
• the matrix summation $e = \text{sum}(\mathbf{C})$ is 38 µsec (ndarray is 84 µsec).

For the transposed elementwise multiplication problem, RSTSR's performance improvement is greater than the 16-times ideal parallel efficiency. As a bandwidth-bottleneck problem, normal parallel performance improvement depends on the specific problem, ranging 0.8–8 times (for simple computation problems with good contiguity, parallelism is sometimes a negative optimization). This may be related to RSTSR's iterator design.

For RSTSR, we estimate the above RI-MP2 energy summation problem takes 21 sec, about 9 times faster than NumPy, much smaller than the 328 sec of the computational bottleneck of the MO-basis 4c-2e ERIs. This is how the problem is solved: for the parts that are not performance-critical, even if the most convenient implementation is used rather than the most efficient algorithm, the program efficiency must still be satisfactory. But it should be pointed out that PyTorch may have implementation efficiency close to or better than RSTSR.

7.4. What Design Flaws Does RSTSR Have That May Need Early Discussion or Complete Refactoring?​

Under the current RSTSR framework, I have not come up with good solutions for some problems. I also hope that the following problems I have thought of, and the problems I may not have thought of, will be discussed as early as possible:

• Complex matrix conjugation. Currently RSTSR's implementation is consistent with NumPy, i.e., a.conj() produces a new matrix with memory reallocated. Therefore, computations like a.conj().T @ &b are inherently inefficient and memory-consuming. This can of course be solved by introducing safe BLAS interfaces, letting users choose whether to use more efficient BLAS functions; but this creates some syntactic noise.
• Matrix types. We do not always deal with arbitrary dense matrices; we may also deal with diagonal matrices, tridiagonal matrices, symmetric matrices, antisymmetric matrices, Hermitian matrices, anti-Hermitian matrices, lower triangular matrices, strictly lower triangular matrices, upper triangular matrices, strictly upper triangular matrices, banded sparse matrices, etc. As a tensor library, we may not need very complete support for two-dimensional matrices; NumPy is also somewhat criticized in this respect, and later GoNum, nalgebra, Faer implemented some of these matrix types differently from ordinary matrices. But if we decide to support these types of matrices, it means we need to re-discuss the storage and computation logic of tensors.
  • At the same time, it should be pointed out that GoNum, nalgebra, Faer are all clearly libraries that only handle two-dimensional matrices or one-dimensional vectors. I have not found an n-dimensional tensor library that supports different matrix types.
  • It is also added here that PySCF generally decompresses lower triangular packed matrices (row-major) before computation. This approach usually does not introduce too much extra efficiency loss.
• Automatic differentiation. To be honest, this is limited by my personal ability. I have not yet grasped the principles of computational graphs. Electronic structure certainly does not need automatic differentiation, but this is now the basic requirement of the AI direction for math libraries. We are of course not unable to make an inference-only library; implementing a few operators without derivatives should not be difficult, but that would be boring. If we want to lean toward the AI direction, I think experienced people need to get involved; but in the end, I myself work on electronic structure, and indeed have no motivation to do automatic differentiation. I do not know how big the demand for AI + first-principles is, nor whether candle and burn can now serve as automatic-differentiation tensor libraries in Rust and be applied to chemistry problems.
• Compilation time is too long. I do not know the cause or solution. It may be that there are too many impls in one type.

The following are problems I have the ability to solve, but they also need early discussion:

• Col-major. The current RSTSR defaults to row-major. In RSTSR, the concrete algorithms are all implemented column-major; column-major indeed has better efficiency. But we completely separated the concrete algorithms from the high-level interface; currently the high-level interface is only designed and implemented in the row-major environment. Whether we should explicitly support col-major and implement its corresponding broadcasting rules for col-major, I am confident I can implement, but if there is a need, it needs early discussion.

  • We need to explain that when RSTSR internally implements tensor addition and subtraction, it will try to transpose tensors to col-major before computation. Therefore, whether the high-level interface is col-major or row-major does not affect computational efficiency. Tensor transposition certainly has a performance cost, but this cost is nothing more than the addition, subtraction and multiplication of five or six integers (layout manipulation), which is insignificant in electronic structure problems.
  • The broadcasting of col-major should be redesigned. NumPy's broadcasting rules are clearly more friendly to row-major. I do not know whether other tensor libraries have made similar attempts.

• Data types and cast principles. For example, np.log(2) gives a floating point number, but the current crate num does not provide a log function for integer types or their corresponding traits. Which functions, which traits, and how to implement them, these problems need further discussion. Electronic structure generally does not use common and special functions; even if needed, although not very elegant, we have ways like .mapv(|x| x.log()). The Map function design is modeled after ndarray.

• Complex trait impls. The current RSTSR seems to have trait impls that are too tedious on some problems, so much so that the Rust compiler can no longer automatically infer types. I encountered this when writing RI-CCSD, and it is one of the least smooth parts of implementing RI-CCSD in my view. For example,

  let c = 2.0 * &a.slice(3);
  

  In the RI-CCSD implementation, this may very likely fail to compile because the type cannot be inferred. You must add this line to make it compile:

  let c: Tensor<f64, _> = 2.0 * &a.slice(3);
  

  In my Rust development, except for the collect::<Vec<T>>() function, I have never encountered a similar problem. I think we need to find a way to avoid errors caused by overly complex impls.

8. RSTSR Design​

8.1. What Is RSTSR's Data Structure? Is It Suitable for Any Dense Tensor?​

RSTSR's current data structure is shown in the figure below:

I believe the above figure is already intuitive, so it will not be expanded in detail here.

RSTSR's code has undergone a complete refactoring, and it differs from the initial data structure. Please refer to pages 12–13 of the RSTSR 1st report. Specifically, the current data is under storage; previously it was the other way around.

RSTSR's data structure is not suitable for all types of dense tensors, especially small dense tensors. In RSTSR's tensors, besides tensor.data which stores the raw data, shape, stride, offset, device also need to be stored. We can call these the tensor's metadata.

• shape is Vec<usize> or [usize; N], its length depends on the tensor dimension; generally no more than 512 bits;
• stride is the same as shape, but the type is isize; generally no more than 512 bits;
• offset is usize, 8 bits in length;
• device depends on the concrete backend:
  • DeviceCpuSerial is just a marker and does not occupy stack space;
  • backends involving parallelism need to use DeviceCpuRayon; it consists of the CPU thread count, the Arc of the global thread pool, and the Arc of the single-thread pool, totaling 24 bits of stack space; creating the global thread pool and the single-thread pool is estimated to need at least 2048 bits of heap space, but Arc pointers guarantee that all tensors share the same heap space, so this heap space can be counted outside the tensor.
• Therefore, a TensorBase type needs an extra 1024 bits, i.e., 16 floating point numbers, besides the data.

For electronic structure, when the basis set size exceeds 100, the redundant information needed to store a tensor in memory does not exceed 2‰, which should be acceptable. But for game rendering, one dimension is generally fixed to 3, and the other dimension is also often fixed and no more than 6; in this case, fixed dimensions are better, i.e., increasing the size of the compiled binary by dispatch (specializing) matrix types of different fixed dimensions (to represent the tensor's metadata), while only raw data is needed when storing the tensor. Libraries like nalgebra, dfdx support fixed dimensions and are more suitable for the latter work. But general scientific computing handles very large matrices and does not need fixed-dimension support.

I have some misgivings about why machine learning libraries support fixed dimensions. For the earliest MLPs and CNNs, models generally fix the hidden layer or channel dimensions; CNN convolution kernels are generally $3 \times 3$ or $5 \times 5$ in dimension, and also have the feature of small individual dimensions (but at the same time it should be noted that CNN convolution kernel parameters also include input and output channels; if both are 64, the tensor size is at least $3 \times 3 \times 64 \times 64 \approx 35000$, which is definitely not a small tensor). Fixing some dimensions helps compiler optimization, and also helps JIT optimization; but in the limit, these automatic optimizations either cannot beat manual optimization, or can judge dimensions and implement with dispatch in a dynamic-dimension framework.

8.2. What Are the Tensor Layout Operations? Do They Cost Computational Performance?​

Typical Layout operations include

• Basic indexing, i.e., extracting sub-tensors or sub-matrices according to certain rules; but depending on the situation, it may also add several dimensions;
• Transpose, i.e., the rotation of several index labels of a tensor;
• Reshape, i.e., changing the shape of a tensor.

Among these Layout operations, basic indexing and transpose do not concretely change the tensor's data. For reshape, when the tensor's underlying data has a certain contiguity (which is generally the case), it also does not change the tensor's data; but when the tensor's data is not contiguously stored, it depends on the situation.

When the tensor's data is not changed, layout operations cost no computational performance. If there is a cost, it is often only a complex operation within 10 integers; compared to electronic structure problems, this amount of operation is unimportant. We need to emphasize that what we do is electronic structure, because I estimate that in some fields, people do care about this little bit of cost.

In the RSTSR program, apart from designing layout-based iterators, we do not optimize the performance of layout operations at all (including but not limited to broadcast, split_at, translate_to_col_major, etc.). Clone when it is time to clone, and reach for Vec<usize> on the heap when it is time, without any hesitation.

8.3. How Are Tensor Iterators Designed?​

First, the tensor iterators we implement are all based on layout iterators. How a tensor is iterated depends on the tensor's shape, stride, offset, and has nothing to do with the tensor's elements. Therefore, the tensor's iterators and the layout's iterators should be separated.

Layout iterators iterate over the tensor's indices (for example, for a matrix it is the two-dimensional usize array [i, j]), and output the relative address in memory of that index as usize. Based on this, the tensor's iterators (which also iterate over indices, but output values or their references) are easy to implement.

In the file layout/iterator.rs, we implemented the layout iterators. Some of the key points are

• Both col-major and row-major versions are written. Layout iteration is not iteration of a one-dimensional array; it can have different orders. For col-major iteration, since the first index changes most, the efficiency is somewhat higher; but this has not yet been proven with micro benchmarks. In the implementation of internal operators, we basically always use the col-major version.
• The most performance-critical functions are next_iter_index and back_iter_index. They manually unroll 1-D – 4-D layouts, and use loops for the rest.
• The RSTSR library allows fixed and variable dimensions. Iterators are much more efficient under fixed dimensions, possibly because under static dimensions shape is implemented with [usize; N], while under dynamic dimensions it is implemented with Vec<usize>. At compile time, [usize; N] is friendlier to the compiler's performance optimization.
• Iterators are frequent operations and need high efficiency; in next_iter_index and back_iter_index, the changes to the internal variables index_start and index_end of the iterator need to be in-place, and cannot clone one out.

In the file feature_rayon/layout_par_iter.rs, we implemented the parallel layout iterators. This is basically just applying the formula. Note that

• I guess that only by implementing IndexedParallelIterator can Rayon's parallel efficiency approach OpenMP.

8.4. How Are Elementwise Tensor Operators Implemented?​

Elementwise tensor operations include the four arithmetic operations and common functions (like sqrt, sin, etc.).

For the four arithmetic operations, they are mentioned in section 3.6 (pp 15–20) of the RSTSR 1st report.

Here we take another, more intuitive sin function implementation as an example.

• RSTSR's operators separate frontend and backend as much as possible; specifically, the Tensor trait part, the Device trait part, and the concrete Device impl part.

• RSTSR's trait definitions generally have the API suffix, to distinguish them from struct types.

• The implementation of the sin function in the Tensor trait part is in tensor/operators/op_unary_common.rs. The sin function takes a single-variable input, so it is called a unary function. Since we need to implement many unary functions (like sqrt, log, etc.), macro_rules needs to be used. The trait in the Tensor part should guarantee that any backend can execute it; that is, it really only designs the interface and simple execution logic here, and the real computation should not appear here.

• The Device trait part is a pure interface. The trait declaration of the sin function is in storage/operators/op_binary_common.rs. Because the sin function involves an input tensor and an output tensor, there are two parameters involved; we call it a binary function (this naming convention is probably not good, and can be changed later).

• Taking DeviceCpuSerial as an example, the Device impl is in device_cpu_serial/operators/op_binary_common.rs. We notice that what is actually done here is

  impl<...> DeviceSinAPI<...> for DeviceCpuSerial {
      fn op_muta_refb(...) {
          self.op_muta_refb_func(..., |a, b| *a = b.sin());
      }
  }
  

  This involves a function op_muta_refb_func, which will go to the code op_muta_refb_func_cpu_serial that actually does the computation. This function reads the input tensor's raw data b: &[TB] and layout lb: &Layout<D>, the output tensor's raw data a: &mut [TA] and layout la: &Layout<D>, and the mapping function f: impl FnMut between them (in the current example, the sin function). This is done to simplify the code, because unary functions all have a unified implementation strategy.

• In the concrete computation code, we call the function translate_to_col_major. Its purpose is,

  • if the input is a col-major tensor, we iterate in col-major;
  • if the input is a tensor of another major, then we transpose it to col-major for iteration. Users only care that the output is correct, so we do not necessarily have to follow the tensor's original iteration order.

• Later, translate_to_col_major_with_contig is called once more, whose purpose is

  • to extract the largest contiguous dimension length. When a tensor has at least one dimension contiguous in memory, we merge all these contiguous dimensions together, and pack the remaining non-contiguous dimensions into one block. For example,
    • if the 3-d tensor a is declared in contiguous memory, then the largest contiguous dimension length is the whole tensor size, with no remaining non-contiguous dimensions;
    • for the above tensor a, taking the slice a[1:4, :, :], it is globally contiguous just like before;
    • taking the slice a[:, 1:4, :], it has one non-contiguous dimension, and the largest contiguous dimension length is the product of the last 2 dimensions;
    • taking the slice a[::-1, ::-1, ::-1], it is reversely contiguous; in this case we consider it completely non-contiguous.
  • when memory is contiguous, using the standard library's built-in contiguous iterator is the fastest; for simple four-arithmetic binary problems, the contiguous parallel iterator may not be as efficient as the contiguous serial iterator; using layout iterators only drags down the speed.
  • the remaining part uses layout iterators.

In summary,

• The Tensor and Device traits are used to define operations; this is the core design of the tensor library (but other users should also be able to implement them, and in the future these operations may be separated from the rstsr-core core library); these parts are not responsible for computation, only for interfaces.
• The Device impl part can be done by users themselves; they can do any high-performance implementation inside.

This idea can be considered modeled after the burn library, but no reference was made to their code.

8.5. Are Transpose and Reshape Operators?​

No. They are counted as Tensor manipulation.

Not all tensor operations are classified as operators.

• Some cases are combined operators, like exp(x) - 1. Depending on the situation, this combined operator is implemented as a standalone operator (for higher floating point precision or higher performance), or can be naively implemented as exp(x) - 1.
• Transpose is indeed not an operator. A Tensor is divided into the concrete data part (storage) and the layout; transpose only handles the layout. The interface at the Tensor level is defined in the into_transpose_f function, while the program that concretely does the layout-level transpose is Layout::transpose.
• Reshape is indeed somewhat special, because it may not need to allocate new memory to store the tensor, but may also need to. Following ndarray's approach, we implement the TensorCow type; if new memory indeed needs to be allocated, we need to use the operator empty (empty_impl) to declare a block of memory, and use the assignment operator (assign_arbitary) to copy the original tensor into the new tensor. The reshape function is actually implemented in the program change_shape_f; but on the one hand this function can still be optimized, and on the other hand it seems too complex in logic and should be separated into the layout part.

8.6. How to Control the Number of BLAS Threads?​

RSTSR learned from rest_tensors in this respect; that is, openblas_set_num_threads is used to control the number of threads:

• when in a Rayon parallel region (when rayon::current_thread_index returns a value), single-threaded BLAS is used;
• when in a serial region (when rayon::current_thread_index returns None), multi-threaded BLAS is restored;

This may not be the best solution, but it should be enough for us.

Not the most urgent problem, but doing this has other detailed issues:

• the above implementation strategy seems to work only with OpenBLAS compiled with pthread; it causes problems with OpenBLAS compiled with OpenMP;
• this function is OpenBLAS-specific and cannot be used when linking MKL or BLIS. Using this function means binding to OpenBLAS.

I think in C/C++, this kind of problem can be solved by defining precompilation macros in header files (like #if(HAVE_OPENBLAS)), and using precompilation judgments to switch between MKL and OpenBLAS in actual use. Such an approach is not impossible in Rust, but I think there is another way: design two sets of backends, DeviceOpenBLAS and DeviceMKL. The two sets of backends use two different thread control strategies. In electronic structure programs, then define through precompilation macros (cargo features in Rust) whether the default DeviceBLAS is DeviceOpenBLAS or DeviceMKL.

Abstract​

When writing scientific computing programs, the data storage structure and how to use that data structure for efficient computation are important and fundamental problems that need to be planned early. Today, NumPy and PyTorch have almost become the de facto standards for high-dimensional tensor data structures; the latter also provides a relatively unified interface for heterogeneous programming. But outside Python, except for the LibTorch library in C++, few languages support both high-dimensional tensors and heterogeneous programming. We hope the RSTSR program can resolve this dilemma in the Rust language in the future. The current RSTSR program is based on the Python Array API standard (a subset of the NumPy API), and has correctly implemented some important high-dimensional tensor operations and manipulations, exposing part of the interface in preparation for high-performance CPU computation and GPU computation. But for the goal of reimplementing most computational chemistry concerns on CPU, the current progress still needs 5–10 weeks; CUDA support may need even more time and effort. This document will 1) briefly analyze the current state of existing matrix or tensor libraries; 2) explain the practical significance and goals of the Rust tensor library RSTSR; 3) provide technical details on the features already implemented in RSTSR.

1. Analysis of the Current State of Matrix or Tensor Libraries​

1.1. Matrix or Tensor Libraries Used by Computational Chemistry Programs​

We write tensor libraries, which could be a contribution to the scientific computing community, or even to machine learning and visual rendering communities; but in the final analysis, our goal is still to serve computational chemistry programs. To that end, let us briefly review the matrix or tensor libraries used in computational chemistry.

This table only shows that a rough program survey was carried out, and does not necessarily prove anything:

• Q-Chem spans a very long development history, developed from the 95's to today, with many old but still usable programs (e.g., ccman), and also libtensor, a fully-featured tensor library developed in modern C++ (e.g., ccman2); different programs use completely different tensor or linear algebra libraries, with no unified programming standard and a large variance in code quality.
• PySCF did not deeply optimize for efficiency and large-system algorithms at its design time, but can achieve efficient implementations on small systems. This can be done through Python/C bindings, splitting performance-critical bottlenecks out of Python into C implementations. In fact, in my RI-MP2 polarizability calculation work, I used a similar technique.
• Psi4 put considerable effort into NumPy performance improvements and tensor contraction standards, but this work was not really brought into Psi4's C++ main program; instead it derived from the Psi4NumPy teaching project. Psi4 is indeed very efficient, but it has no unified tensor library; rather, it handles tensor contraction problems separately in specific tasks (e.g., occ).
• MPQC is probably a computational chemistry program that has in fact ended maintenance; it is not a successful project. But E. Valeev is still actively developing libint and TiledArray.
• CP2K once independently developed libsmm, but later gave way to Intel's libxsmm. The latter can be used not only for DFT calculations, but is also an important implementation tool for CNN-representative machine learning methods on CPU.
• There are many programs I do not know well, especially those that use MPI at scale (e.g., FHI-Aims, ORCA, VASP) or GPUs (e.g., TerraChem).

But the commonality among them is,

• almost all successful modern computational chemistry programs have strong support for tensor computation.

1.2. PySCF's Usage Characteristics and Their Inspiration for Rust Program Development​

Every computational chemistry program has its own character; no computational chemistry program truly agrees on how to use math libraries (of course, disagreement also covers program structure and interfaces, input/output file formats, algorithms, etc.).

Perhaps the most successful molecular computational chemistry program today is ORCA, but this may be attributed to

• its forum-style community maintenance;
• its complete program manual and relatively few program bugs;
• its functionality that is relatively complete compared to Gaussian, with considerable performance improvement;
• its free usage and convenient installation strategy, and multi-platform availability.

But from a developer's perspective, ORCA is far from a good choice, especially due to its closed-source nature. Among open-source programs, the most successful today is PySCF. Its success is of course partly due to the PySCF developers' own characteristics:

• a clear strategy of separating code logic (Python) from performance-critical code (C);
• usability-oriented design (not performance-oriented or system-size-capacity-oriented);
• a breakthrough in designing the computational chemistry program as a library rather than a main-program-driven control flow; this is crucial for computational chemistry method developers, and also convenient for users with high customization needs in computation processes;
• flexible use of object-oriented and functional programming features, prioritizing interface usability and avoiding complex inheritance;

But many of its features are not the result of the PySCF developers' work, but rather of their choices:

• NumPy is almost the standard for medium-scale dense tensor computation (larger than the matrices in machine learning, game engines, and rendering tasks, but not requiring cross-node parallelism), which lowers developers' learning cost, and its code is easy to write and highly readable;
• Python as a scripting language can run and debug immediately, without spending lots of time compiling and troubleshooting.

When we choose Rust for development, we inevitably give up some of PySCF's advantages, especially Python's advantages as a scripting language; but what we hope to gain (or can gain) in exchange is

• program execution efficiency;
• program stability and lower memory overhead (avoiding memory leaks caused by programming or the language itself);
• binaries that are easier to distribute;
• less FFI (cross-language interaction), and more convenient cross-platform (cross-OS) support;

Relative to C++ and other languages, Rust also has greater advantages in engineering standards and deployment:

• Cargo's convenient build, test, rustfmt style, clippy code standards, tarpaulin coverage, and precompilation feature options; though documentation (doc) may be Cargo's current weakness;
• thanks to Rust's strict code standards, as long as compilation produces no errors, we can safely handle complex lifetime and trait combinations.

Correspondingly, the advantages of other languages that Rust inevitably gives up (which may also be disadvantages) include:

• Rust is not a scripting language; whether or not there is a Jupyter development environment (evcxr), this is definitely inconvenient;
• Rust does not support template specialization as freely as C++, which can standardize code, but also makes lazy evaluation and type inference requiring compile-time computation difficult.

1.3. NumPy's Characteristics and Their Relevance to Computational Chemistry​

In my view, PySCF's success is not only due to its developers' own efforts, but also to riding the tailwind of Python and NumPy. We can use concrete examples to show NumPy's applications in computational chemistry; and thereby reflect on what expectations, as program developers, we might actually have of math libraries.

1. High-dimensional tensors

   These are frequently used in computational chemistry, especially in algorithms involving CC methods.

2. Broadcasting

   $$\Delta_{ij}^{ab} = \epsilon_i + \epsilon_j - \epsilon_a - \epsilon_b$$

   D = (
     + eo[:, None, None, None] + eo[None, :, None, None]
     - ev[None, None, :, None] - ev[None, None, None, :]
   )
   

   This code may not actually be used in real MP2 energy calculations, because it has large memory requirements. But MP2 can be implemented in Python in three lines of code; when performance is not the critical part, such code is very intuitive. Similarly, in Laplace-Transform OS-MP2, we also use multiplicative Broadcasting:

   $$\mathcal{G}_{gia} = e^{D_i^a t_g}$$

   grid_exp = np.exp(D_ia * grid_points[:, None, None])
   

   Similar situations are also common in DFT grid integration.

3. Dimensional transformation under contiguous memory

   Although the following problems can be implemented with Einstein summation / einops, einops should in principle satisfy the following computational requirements. But efficient einops implementations are difficult and require substantial effort; without einops, the best approach is to reshape the dimensions and then do matrix multiplication:

   $$\begin{aligned} D_{jk}^{(2)} &= \sum_{bia} -2 T_{jb,ia} t_{kb,ia} \\ D_{ab}^{(2)} &= \sum_{jci} 2 T_{jc,ia} t_{jc,ib} \\ \Gamma_{P,ia} &= \sum_{jb} Y_{P,jb} T_{jb,ia} \end{aligned}$$

   T, t = [tsr.reshape(nocc, nocc * nvir**2) for tsr in (T, t)]
   rdm2[so, so] = -2 * T @ t.T
   T, t = [tsr.reshape(nocc**2 * nvir, nvir) for tsr in (T, t)]
   rdm2[sv, sv] =  2 * T.T @ T
   T = T.reshape(nocc * nvir, nocc * nvir)
   Gamma = Y @ T
   

   The np.reshape above needs to transform the 4-d tensors into 2-d matrices, and should not involve memory copies.

4. Einstein summation / einops

   Still for the above problems, using einops would be much more convenient:

   rdm2[so, so] = -2 * np.einsum("jbia, kbia -> jk", T, t)
   rdm2[sv, sv] =  2 * np.einsum("jcia, jcib -> ab", T, t)
   Gamma = np.einsum("Pjb, jbia -> Pia", Y, T)
   

   Although the above code is MP2 gradient-related code, similar code is common in the CC computation modules of MPQC, Q-Chem, and PySCF.

   For example Q-Chem:

   // ccman2/ccman2/cs_cc/cs_ccsd_pt.C
   contract(d, t2a_re(c|d|i|j), i_vvov_re(a|b|k|d))
   

   For example MPQC:

   // ccsd.h, line 342
   tau("a,b,i,j") = t2("a,b,i,j") + t1("a,i") * t1("b,j");
   

   Even code this intuitive, if combined with an efficient einops, can reach considerable efficiency; but einops is not an easy library to write.

NumPy also has some difficult-to-handle problems; in PySCF, some computation processes must rely on C implementations, and problems that are difficult to implement efficiently in Python unless using numba, including

1. matrix transposition involving memory copies;
2. manual parallelization and triangular matrix multiplication (e.g., AO2MO computation).

Some of these problems are inherent to NumPy itself, and some are Python language problems.

Overall, in computational chemistry, we need matrix multiplication functionality, but we also have substantial needs for the high-dimensional tensor features supported by NumPy and its distinctive features.

1.4. The Python Array API and Existing Matrix/Tensor Libraries​

NumPy not only interfaces with the functionality needed by computational chemistry, but the convenience of its API is also crucial. In the 10's, NumPy had already become Python's de facto numerical computation program; during the explosive growth of machine learning in the 15'–20's, amid the development and competition of TensorFlow, MXNet, PyTorch and other libraries, in my view, PyTorch's current success is largely because its API is very close to NumPy, making it easy for beginners, avoiding high learning costs, and reducing the communication cost between libraries.

The authors of Python math libraries also realized the importance of API interfaces. In 2020, a community group discussed the interface forms of Python math libraries and, based on NumPy and other popular math libraries, defined the API interfaces that math libraries must satisfy, in the form of the Python Array API Standard (after NumPy 2.0, the Python Array API Standard is a true subset of the NumPy API).

Although the Python Array API is a Python interface, other languages were also considered in its design; C++ and Rust, among others, can also follow the rules of this interface to a certain extent for program development. I personally also hope that the library we develop will take the Python Array API into account as much as possible.

1.4.1. The Python Array API and Its Related Math Libraries​

Currently, the vast majority of successful math libraries have APIs close to NumPy, or are themselves libraries designed under NumPy's influence and inspiration.

But it should also be pointed out that successful math libraries do not necessarily follow the Python Array API; this is also reflected in the Rust language, which we will supplement in the next subsection. For other languages,

• Fortran, Matlab, Julia and other languages natively support high-dimensional tensors, and partially support some math operations;
• many Fortran, C/C++ programs directly use high-performance BLAS; not necessarily using external wrapper libraries or self-developed simple wrappers;
• Eigen in C++ is somewhat special: it has strong lazy evaluation support and small-matrix computation performance; large-matrix performance is not bad. But it was developed in the 05's, developed early, and is hard to turn around.

But in general, today's math library designs can usually express tensor types, and at the API level try to be as close to NumPy (or the Python array API standard) as possible.

1.4.2. Math Libraries in the Rust Language​

The math and machine learning libraries in the Rust ecosystem are quite diverse, but not all satisfy our needs. The table below briefly lists them.

1. ndarray is not expected to implement GPU support (ndarray #1377).
2. ndarray's linear algebra is provided by the external library ndarray-linalg; the latter has stopped maintenance. But conversions between ndarray and faer, nalgebra are relatively easy; therefore, despite the fragmented ecosystem, solutions that handle both high-dimensional tensors and linear algebra still exist.
3. faer has strong complex number support, but it is essentially not a general matrix library but a high-performance computing library benchmarking against OpenBLAS, so it only implements common floating point types. The library author is aware of F16 and BF16, but they will not be introduced into faer in the short term (faer-rs #32).
4. faer currently supports fewer linear algebra functions, but has high performance.
5. burn does have high-dimensional tensors, but this is mainly for storing the tensors needed by machine learning. The library itself does not support strides, so tensors cannot be used for general computation.

Besides these, einops, as a supporting library, also has important potential help for computational chemistry programs. The closest libraries in Rust are einops (serving tch-rs) and candle_einops (serving candle), i.e., these libraries mainly target machine learning applications.

2. Significance and Goals of the Rust Tensor Library RSTSR​

From the above discussion, I believe that the current Rust math libraries still lack something for today's computational chemistry programs. Developing a new math library may not be valuable in the short term, but it is necessary in the long term.

For this reason, I have tried to start the RSTSR program, and hope to invest time in this project. I hope such a program will aim to assist chemistry program development, while also accommodating existing tensor library standards and other possible scientific computing needs to a certain extent.

2.1. Breakdown of Tensor Library Features​

For computational chemistry, I think the functionality a math library must carry includes at least

1. high-dimensional tensor data structures and their basic operations (Python array API standard);
2. basic linear algebra (matrix decomposition, eigenvalue problems) and FFT (if periodic systems are supported);
3. multithreading or other parallel modes (guaranteeing that at least 30% or more of the machine's floating point or bandwidth efficiency can be used);
4. complex number types, and preferably arbitrary types (especially arbitrary floating point precision types);

The above characteristics are also (a subset of) NumPy's functionality. Many of these features also appear in the Rust library ndarray, but it is weak in linear algebra, and lacks the following important optional features.

The important optional features include:

5. advanced linear algebra (interpolation, quadratic convergence and solving, matrix function extrema, though not necessarily as a standalone math library);
6. einops (simplifying tensor computation code and unifying tensor multiplication implementations, but replaceable with basic matrix algebra);
7. GPU heterogeneity (devices with better cost-performance and faster computation);
8. special BLAS features (batched GEMM may be used for DMRG methods, BF16-based FP32 GEMM may be used for post-HF methods);

Features of uncertain importance include:

9. MPI heterogeneity (depending on the technical path, it may be Dask-native support or Scalapack non-native support);
   Whether MPI matters depends on 1) whether the program's main parallel mode uses threads or processes, 2) whether the program focuses on large-scale parallelism; many computational chemistry programs have little multi-process support, yet are very successful;
10. small matrix multiplication (some DFT computation requirements);
11. symmetric tensor operations, tensor symmetrization/antisymmetrization (symmetry systems, CC algorithms);
12. general sparse matrix data structures and basic operations;
13. computational graphs and automatic differentiation (the most important feature of machine learning programs, but computational chemistry itself usually does not need it; if our program needs to interface with machine learning, this feature needs to be considered);

Features that are not necessarily important include:

14. lazy evaluation (it is generally used in memory-bottleneck code, but matrix-multiplication-dominated problems are either not memory-bound, or can sacrifice a small amount of code elegance to guarantee efficiency; lazy evaluation is somewhat difficult to implement, which would bring great trouble to library maintenance and open-source collaboration).

2.2. Goals of the Tensor Library RSTSR and Expected Implementation Phases​

In the tensor library RSTSR, we will focus on implementing high-dimensional tensor data structures and GPU heterogeneity. The work completed so far concentrates on high-dimensional tensor data structures and their basic operations:

• high-dimensional tensor data structures;
• most Layout operations (tensor operations that do not change the underlying data);
• Layout-based tensor element iterators (iterators decoupled from the underlying data);
• broadcasting (matching rules for tensors of different dimensions);
• for computation problems, separating backend implementations from frontend interfaces;
• basic arithmetic operations on tensors (a fairly efficient single-threaded implementation);
• tensor matrix multiplication (a correctness-first implementation; efficiency needs to be improved through future backends);
• tensor creation and shape changes.

These are preliminary works, and are not yet mature enough to be applied to the REST program.

If I could develop the RSTSR tensor library full-time, I think a suitable schedule would be:

• 2–4 weeks to fully reproduce most of the functionality required by the Python array API (excluding linear algebra);
• 2–4 weeks to complete the parallel computation code, the BLAS/Lapack backend, and the faer backend;
• 2–3 weeks to implement the supplementary features in rest_tensor (including linear algebra);
• 4–8 weeks to implement the CUDA backend based on cudarc;
• 2–4 weeks to propose an MPI solution based on Scalapack (but not written as a library).

In the above process, the first two steps need to be done first; the order of the subsequent steps can be shuffled. The total is about 3–5 months.

3. RSTSR Design and Technical Details of Implemented Features​

The Rust language usually clearly divides composite types into two parts: 1) data structures (struct) or interface traits (trait); 2) their implementations (impl). In general, implementations (impl) are easy to replace; but changes to data structures (struct) and interface traits (trait) easily lead to serious code refactoring.

Among the features listed in the "Breakdown of Tensor Library Features" section above, the ones that profoundly affect data structures and interface traits are:

1. high-dimensional tensor data structures, and their basic operations;
2. GPU heterogeneity;
3. symmetric tensors;
4. computational graphs and automatic differentiation;
5. lazy evaluation.

The other features are relatively independent and will not significantly affect the tone of the program design, except: 9. MPI heterogeneity: whether or not the native-support path is adopted, MPI heterogeneity cannot be implemented directly on a standalone high-dimensional tensor; it will inevitably split a single large high-dimensional tensor into small pieces; therefore, MPI heterogeneity must be built on top of well-developed high-dimensional tensor data; 12. sparse matrices: they differ too much from high-dimensional tensors, and would generally use a data structure different from that of high-dimensional tensors.

RSTSR currently only considers high-dimensional tensors and future GPU programs. Although we call GPU "heterogeneous" here, we currently only regard it as a backend different from CPU. In fact, when we also treat BLAS and faer as two tensor computation backends, then high-performance CPU libraries themselves are backends, implemented in the same way as future GPU implementations. Therefore, multi-backend implementation should also be considered early.

3.1. Abandoned Feature: Automatic Differentiation and Lazy Evaluation​

Automatic differentiation requires computational graph functionality. This type of feature has at least two core difficulties:

• expression trees;
• variable ownership.

Expression trees are somewhat difficult to implement. They are essentially lazy evaluation as well, and can implement asynchronous computation (CUDA stream computation) within this framework, and can also simplify computation flows (e.g., c += 2 * a simplified to famdd(c, 2, a)). But they are not a common data structure in scientific computing. Implementing their inplace operation rules is difficult; even PyTorch frequently had inplace operation bugs in its early days.

At the same time, as Rust is a language with very strict variable lifetimes, tree-structured implementations are indeed difficult. In fact, when implementing automatic differentiation, candle and burn directly use RwLock or Arc smart pointers, avoiding ordinary lifetime-affected variables to represent tensors.

My view is that if we need to introduce automatic differentiation in the future, we need to do the following two things:

• write a separate AI frontend program that uses RSTSR as a backend, rather than writing a frontend directly in RSTSR;
• seek cooperation with technology companies.

Developing a self-made automatic differentiation program without connecting to industry needs is not necessarily meaningless, but it is difficult to win an audience; assuming Rust can indeed do machine learning, then in the end, projects led by technology companies will be more widely adopted than ours, so that we would waste too much time before our programs are abandoned. If our goal is only multi-backend scientific computing, then we have much more freedom in programming, and the program difficulty is greatly reduced; before large companies do this at scale, I am confident that we can do it well, at least without wasting time.

3.2. Abandoned Feature: Overloading Assignment Statements​

This is a Rust feature. In C++, the equals sign = can be overloaded with the operator= function; Rust does not allow this.

The result is that at least the following two convenient C++ practices are infeasible in Rust:

• regarding lazy evaluation, if we care about memory reuse, a typical example in C++'s Eigen library is

  mat1.noalias() = mat2 * mat2;
  

  This guarantees that mat1's memory can be reused, but this is achieved through operator= overloading.

• similar to MPQC's tensor contraction

  tau("a,b,i,j") = t2("a,b,i,j") + t1("a,i") * t1("b,j");
  

  the contraction operation needs the index-label information on both sides of operator=, and cannot be simply completed by assigning the RHS to the LHS.

3.3. RSTSR Data Structures​

The RSTSR project has learned (or plans to learn) strategies from many other libraries:

• ndarray: data structures, lifetime management, vectorization of reduce operations;
• candle + cudarc: GPU backend integration into Rust programs;
• burn: the external API of multi-backend implementations;

As the basic data structure, RSTSR's tensor will be expressed in the following way:

• tensors are split into underlying data (DataOwned) and layout (Layout);
• shapes (Layout) are split into dimension (shape), stride (stride), and offset (offset);
• the dimension type is specified by DimAPI, which can be fixed-dimension arrays ([usize; N]) or variable-dimension arrays (Vec<usize>); efficient computation generally prefers the former.
• underlying data is split into backend (DeviceAPI), underlying data type (RawVec), lifetime and ownership;
• since the backend and the underlying data are somewhat coupled, the split strategy is shown in the figure above.

The above data structure is quite different from both burn and ndarray:

• burn directly defines Tensor<B, D, T>, but it completely relies on RwLock for variable ownership and lifetimes; and burn only supports fixed dimensions;
• our implementation strategy is closer to ndarray; but ndarray only implements on CPU, using a custom data type (equivalent to manually reimplementing Vec<T>, with a fair amount of unsafe code). In RSTSR's implementation, for the CPU backend, the simpler Vec<T> is used as the basic data storage format.

3.4. On the Implementation of Views​

High-dimensional tensors inevitably involve the concept of view. In our implementation, the TensorBase type is defined as

pub struct TensorBase<R, D>
where D: DimAPI,
{
    pub(crate) data: R, // Vec<T>/CudaSlice<T> (with lifetime and backend)
    pub(crate) layout: Layout<D>, // {shape, stride, offset}
}

The view (struct DataRef) is directly defined as

pub enum DataRef<'a, S> {
    TrueRef(&'a S),
    ManuallyDropOwned(ManuallyDrop<S>),
}
pub type TensorView<'a, T, D, B> = TensorBase<DataRef<'a, Storage<T, B>>, D>;

and we generally only use TrueRef, which means our views are literally, in the literal sense, references to the Storage<T, B> type. ManuallyDropOwned is only used to initialize TensorView<T, Ix1> from &'a [T].

This differs considerably from ndarray; they store all the underlying information in ArrayBase:

pub struct ArrayBase<S, D>
where S: RawData
{
    data: S, // Customized Vec<T> if OwnedRepr
    ptr: std::ptr::NonNull<S::Elem>, // offset
    dim: D,                          // shape
    strides: D,                      // stride
}

But if this vector is referenced (ViewRepr), then ArrayBase.data only has a lifetime:

pub struct ViewRepr<A> { life: PhantomData<A> }

We did not adopt ndarray's approach. This is because outside the CPU device, the purpose of a pointer can only be to access the underlying data and its associated backend device information, and should not be used for pointer arithmetic (for concrete tensor computations). This also requires that referenced data also carry the underlying data and its associated backend device information, and cannot be an empty lifetime.

Doing so naturally brings some trouble. The ManuallyDropOwned mentioned above is one of them. But this is much more convenient and intuitive for program maintenance.

3.5. On the Implementation of Layout​

Given the raw data, a tensor is defined by the dimension (shape), stride (stride), and offset (offset):

pub struct Layout<D>
where D: DimBaseAPI,
{
    pub(crate) shape: D,
    pub(crate) stride: D::Stride,
    pub(crate) offset: usize,
    size: usize,  // this may not be necessary and may be removed
}

After clearly separating data from shape, many tensor operations can be explicitly performed only on the shape (layout), without touching the underlying data (data) at all. For example, if we implement a transpose function for Layout<D>, then transposing a tensor TensorBase<R, D> becomes very easy:

// different to actual implementation: `axes: &[I: TryInto<isize> + Copy]`
pub fn transpose<I, R, D>(tensor: TensorBase<R, D>, axes: &[isize]) -> Result<TensorBase<R, D>>
where R: DataAPI, D: DimAPI,
{
    let layout = tensor.layout().transpose(&axes)?;
    unsafe { Ok(TensorBase::new_unchecked(tensor.data, layout)) }
}

While in ndarray, since the layout and the tensor are directly bound, the permute_axes function (equivalent to transpose) must be implemented directly on the tensor type (source of permute_axes).

Operations that only change the Layout without concretely changing the tensor data include at least:

• taking sub-tensors, e.g., a.slice([.., ..3, None, 5..8]);
• transposition, e.g., a.transpose([0, 2, 1]);
• broadcasting (e.g., dimensions [5, 1, 3, 1] and [4, 3, 2] can be broadcast to [5, 4, 3, 2] without creating new tensor data);
• iteration over tensors can be transformed into iteration over Layouts (exporting offset values), making iteration for any backend easy to implement.

3.6. On the Implementation of CPU Backend Tensor Addition​

As an example of separating the tensor library frontend from the backend, we take the addition operation to show the general idea of backend separation, as well as some simple performance optimization strategies.

First, we need to reach a consensus: we do not specially handle small matrix or tensor computations, especially in computational chemistry; therefore, we will assume that the computational cost of tensor operations far exceeds the computational cost of Layout transformations (generally no more than 10 μs).

3.6.1. Function Signatures of Addition and the Backend Separation Strategy​

This part is not a distinctive feature of this library; it has already been implemented in ndarray.

We can notice that addition has several cases:

1. C = &A + &B
2. C = A.view() + B.view()
3. C = A + &B
4. C = A + B
5. C = &A + b (b as scalar)
6. C = A + b (b as scalar)
7. C += &B
8. C += b (b as scalar)
9. ......

There are quite a few cases that can appear here, but in general, addition operations are divided into 6 cases:

1. add_tenary(&mut C, &A, &B),
2. add_binary(&mut A, &B); let C = A,
3. add_assign_binary(&mut C, &B), cases 4–6 are when B is not a tensor but a scalar; we will not discuss these cases for now.

Note that,

• C = &A + &B and C = A.view() + B.view() should be implemented as add_tenary(&mut C, &A, &B);
• C = A + &B and C = A + B should be implemented as add_binary(&mut A, &B); let C = A;
  • when allowed, it will inplace execute A = &A + &B, and then assign C = A;
  • but if the dimensions do not allow it (B.shape() not broadcastable to A.shape()), then execute add_tenary(&mut C, &A, &B);
• C += &B is implemented as add_assign_binary(&mut C, &B).
  • it should be noted that add_binary(&mut C, &B) and add_assign_binary(&mut C, &B) have very similar tasks, but the former is C = &C + &B while the latter is C += &B; + (add) and += (add-assign) are not the same operation. But similarly, both are binary operations with the same function type.

Finally, we summarize the ternary and binary operations as

pub fn op_mutc_refa_refb_func(
    c: &mut TensorBase<RC, DC>,
    a: &TensorBase<RA, DA>,
    b: &TensorBase<RB, DB>,
    f: F,
) -> Result<()> {...}

pub fn op_muta_refb_func(
    a: &mut TensorBase<RA, DA>,
    b: &TensorBase<RB, DB>,
    f: F,
) -> Result<()> {...}

and derive the following assignment binary operation

pub fn op_refa_refb_func(
    a: &TensorBase<RA, DA>,
    b: &TensorBase<RB, DB>,
    f: F,
) -> Result<Tensor<TC, <DA as DimMaxAPI<DB>>::Max, B>> {...}

At this point, the tensor-level function abstraction is complete. The remaining tasks are:

• at the user-usage level, implement the Add, AddAssign traits upward;
• for the backends, implement the concrete tensor addition operations downward.

3.6.2. Function Signatures of the Addition Backend​

The concrete implementation of tensor addition requires at minimum the following data:

• the tensor's raw data (Vec<T> or &[T] on CPU);
• the tensor's shape information (Layout<D>)

In the storage/operators.rs file, for the following cases, we defined operation interfaces:

• add_tenary(&mut C, &A, &B)

    pub trait DeviceAddAPI<TA, TB, TC, D>
    where
        TA: core::ops::Add<TB, Output = TC>,
        D: DimAPI,
        Self: DeviceAPI<TA> + DeviceAPI<TB> + DeviceAPI<TC>,
    {
        fn op_mutc_refa_refb_add(
            &self,
            c: &mut Storage<TC, Self>,
            lc: &Layout<D>,
            a: &Storage<TA, Self>,
            la: &Layout<D>,
            b: &Storage<TB, Self>,
            lb: &Layout<D>,
        ) -> Result<()>;
    }
  

• add_assign_binary(&mut C, &B)

    pub trait DeviceAddAssignAPI<TA, TB, D>
    where
        TA: core::ops::AddAssign<TB>,
        D: DimAPI,
        Self: DeviceAPI<TA> + DeviceAPI<TB>,
    {
        fn op_muta_refb_add_assign(
            &self,
            a: &mut Storage<TA, Self>,
            la: &Layout<D>,
            b: &Storage<TB, Self>,
            lb: &Layout<D>,
        ) -> Result<()>;
    }
  

• add_binary(&mut A, &B); let C = A: the current implementation of this case is lazy; it is currently done in DeviceOp_MutA_RefB_API, without explicit backend separation.

At this point, the frontend/backend separation is complete. For CPU, apply the above traits to DeviceCPU; for CUDA, apply the above traits to DeviceCUDA. CPU has been implemented, but CUDA still has a long way to go.

3.6.3. CPU Tensor Addition Implementation​

Now we describe the CPU implementation of the above trait DeviceAddAPI (or function op_mutc_refa_refb_add). It is implemented in cpu_backend/operators.rs:

impl<...> DeviceOp_MutC_RefA_RefB_API<...> for CpuDevice
where ..., F: FnMut(&mut TC, &TA, &TB),
{
    fn op_mutc_refa_refb_func(&self,
        c: &mut Storage<TC, CpuDevice>, lc: &Layout<D>,
        a:    & Storage<TA, CpuDevice>, la: &Layout<D>,
        b:    & Storage<TB, CpuDevice>, lb: &Layout<D>,
        mut f: F,  // such as `|c, a, b| *c = a.clone() +  b.clone()`
    ) -> Result<()> {
        // re-align layouts
        let layouts_full = translate_to_col_major(&[lc, la, lb])?;
        let layouts_full_ref = layouts_full.iter().collect_vec();
        let (layouts_contig, size_contig) = 
            translate_to_col_major_with_contig(&layouts_full_ref);
        // contiguous if possible, otherwise use iterator of layout
        if size_contig >= CONTIG_SWITCH { // CONTIG_SWITCH ~= 16
            let iter_c = IterLayoutColMajor::new(&layouts_contig[0])?;
            let iter_a = IterLayoutColMajor::new(&layouts_contig[1])?;
            let iter_b = IterLayoutColMajor::new(&layouts_contig[2])?;
            for (idx_c, idx_a, idx_b) in izip!(iter_c, iter_a, iter_b) {
                // compiler should optimize following for loop with SIMD
                for i in 0..size_contig {
                    f(&mut c.rawvec[idx_c + i],
                         & a.rawvec[idx_a + i],
                         & b.rawvec[idx_b + i]
                    );
                }
            }
        } else { // not contiguous after transpose in any cases
            let iter_c = IterLayoutColMajor::new(&layouts_full[0])?;
            let iter_a = IterLayoutColMajor::new(&layouts_full[1])?;
            let iter_b = IterLayoutColMajor::new(&layouts_full[2])?;
            for (idx_c, idx_a, idx_b) in izip!(iter_c, iter_a, iter_b) {
                f(&mut c.rawvec[idx_c],
                     & a.rawvec[idx_a],
                     & b.rawvec[idx_b]
                );
            }
        }
        return Ok(());
    }
}

This is a generic ternary operation implementation; implementing addition through this function is very easy:

op_mutc_refa_refb_func(
    c, lc, a, la, b, lb,
    |c, a, b| *c = a.clone() +  b.clone()
)

Let us briefly analyze the implementation principle of addition. The core part is how to rearrange the layouts of the three vectors. We first assume that the three vectors' layouts lc, la, lb have the same shape.

The core problem of elementwise tensor operations is memory alignment. Relative to the more complex Matrix Multiplication problem, on CPU, elementwise operations are very simple, and it is easy to achieve fairly high performance without complex tricks.

Imagine that there are now three kinds of layouts (assume tensors A, B, C all have the same layout):

• F-contiguous: this is the most convenient case; just scan the memory from head to tail.

• C-contiguous: since our program only implements the F-contiguous iterator (but this is generally enough), we need to transpose it to F-contiguous:

  shape  : [  100, 200, 300] -> [300, 200,   100]
  stride : [60000, 300,   1] -> [  1, 300, 60000]
  

  In this way, the memory of all $100 \times 200 \times 300$ numbers is continuously aligned, and we just scan from head to tail.

• Arbitrary strided with at least one dimension contiguous: this is a relatively special case, where the computation problem can be processed continuously, but is neither C-contiguous nor F-contiguous. Then we transpose the tensor by sorting the strides from small to large:

  shape  : [  100, 200, 300] -> [200, 300,   100]
  stride : [80000,   1, 200] -> [  1, 200, 80000]
  

  After sorting the strides from small to large, another problem arises: the $100 \times 200 \times 300$ numbers are not contiguous: every $200 \times 300 = 60000$ numbers are contiguous, but after that, 20000 data elements must be skipped before reaching the next group of valid data. In this case, if the iterator is used directly, the compiler will not be hinted that the operations can be vectorized (SIMD) optimized. Therefore, we need to find a way to tell the program that this tensor actually has 60000 contiguous numbers; do contiguous computation on the contiguous numbers as much as possible, and jump over the non-contiguous places using the pointer positions given by the iterator.

Finally, it should be noted that if the dimensions of the three tensors do not match, then the addition cannot be vectorized, and must be computed through the relatively inefficient iterator. The inefficiency here does not mean that iterators are bad, but that iterators cannot hint the compiler to vectorize at the -O3 optimization level. For example, adding a C-contiguous matrix to an F-contiguous matrix is itself very unfriendly to memory contiguity. Even so, there are actually better approaches than simply using iterators (even if the contiguous memory cannot be vectorized, there is still a chance to fully exploit the L2 cache), but the implementation complexity would be too high, since we might also have to handle 3-dimensional tensors.

Of course, the current implementation is still fairly fast single-threaded; but this computation problem can also be completed with multithreading. I think we can implement simple, correct tensor addition in the main program (rstsr-core), and create a new parallel backend (DeviceCPURayon) in another library (e.g., rstsr-rayon), to perform parallel (or high-performance) tensor addition. The same is true for matrix multiplication. In addition, the above analysis may also be helpful for efficiently implementing multithreaded tensor addition.

3.7. Proposal of the Matrix Multiplication Symbol %​

I propose that in RSTSR, matrix multiplication (matmul) takes the % symbol. Its implementation has already been completed in tensor/matmul.rs. That is, the following expressions are matrix multiplication:

let c = &a % &b;
let c = a % b;  // in this way, `a` and `b` are consumed

It will have similar functionality to the Python statement

c = a @ b  # c = np.matmul(a, b)

The % symbol is the remainder operator (trait Rem).

Referring to Python's introduction of the @ symbol as matrix multiplication in PEP 465, my considerations when introducing the % symbol into the RSTSR library are:

• replacing matrix multiplication with a binary operator is quite important for code readability; this is explained in detail in PEP 465;
• the reason Rust cannot use the @ symbol is that it is already a pattern binding binary operator; therefore, in the Rust language, it is in any case impossible to use @ for matrix multiplication as Python does;
• the operator precedence of the % symbol in Rust is the same as * and /, i.e., multiplication and division; while in Python, @ has the same operator precedence as *, /, //, %;
• although the % symbol is very common in integer remainder operations, the symbol is almost impossible to use in matrix operations;
• disregarding mirror planes parallel to and perpendicular to the page, both the % symbol and the $\times$ sign have $S_4$ symmetry ($\hat{C}_2 + \hat{i}$).