General block operator scheme

[Slaedr]

The starting point of this idea is that we need a framework for a batch of independent small linear problems. One way of addressing this is to put the small matrices into a large block-diagonal matrix made up of small sparse matrices, where the diagonal blocks may or may not have shared properties such as shared sparsity pattern. We then need to solve the small problems with high performance.

Since we are going the way of treating 'BatchCSR' as essentially a block-diagonal matrix with CSR blocks, does it make sense to re-orient the whole concept, and have an abstraction for block matrices with an arbitrary type of block? For chemical reaction applications with independent reactions (let's call this "CRI" problem for short), we essentially need a BlockMatrix<SmallCsr>, which for this particular application is also block-diagonal. I am sure there are situations where you would want to solve coupled chemical reactions in different cells, for which you need a general BlockMatrix<SmallCsr>, with off-diagonal sparse blocks. This already happens in fully implicit CFD of reactive flows but they sometimes just use dense blocks for smaller chemical systems. This framework is essentially the full generalization of block matrices. This could also be very useful for high-order finite element simulations with variable-dense and sparse element blocks. For the finite element case, non-square off-diagonal blocks need to be supported. Such a framework is now easier since we are moving towards storing everything in global memory. We could have SmallCsr, SmallDense etc., whose main property is that internal data arrays are meant to be stored as views, and they accept pointers to pre-allocated storage areas in the constructor or something. They need not have apply functions, in which case BlockMatrix takes care of applying everything inside. But if the small types have apply functions, we probably want them as __device__ functions. The higher-level data structure of BlockMatrix could just always be block-CSR, possibly. In the beginning we don't need to consider all these generalities, but just creating the required class hierarchy and using a BlockMatrix<SmallCsr> for the CRI problem might make sense. I believe MFEM already has an abstraction for this kind of thing.

We would then also need the concept of BlockPreconditioner<SmallSolver>, where the Block Preconditioner is block-Jacobi, block-Gauss-Seidel, Schur Complement (approximate block LU) etc. The solver type which is the template parameter is used to solve the individual sparse blocks as and when required by the chosen global preconditioner. For the CRI problem, would only need block-Jacobi as the global preconditioner. SmallSolver could be things like SmallPreOnly<SmallParIlu>, SmallGmres<SmallJacobi> etc.; for now we only need something like SmallBiCGStab<SmallJacobi>, say. The Block-Preconditioner could then be used as preconditioner to any of our existing global solvers; for now we only need a preconditioner-only global solver (in PETSc this is KSPPREONLY). Note that all 'Small' things are meant to operate within device kernels, while Block-Preconditioners operate at a global level and launch kernels as needed. SmallSolver could also be called OnDeviceSolver or something.

Further thought is needed on the distributed implementation of these concepts, but it should be very possible. I believe MFEM and PETSc already do this to at least some extent.

[upsj]

I agree that it might make sense to have a common representation for the block structure in batched operators. If we were to add a batched vector LinOp (for which I currently see both advantages and disadvantages, see below), we could even take care of the dimension validation in there. The question whether we need a templated type BlockMatrix to me is mostly base on whether the block matrix part of the LinOp needs to know anything about the CSR matrix. I see the point of having just a single block of storage, though I can imagine that this would could make debugging more complex, e.g. catching out-of-bounds memory accesses when row_ptrs and col_idxs are directly adjacent.

Templates in Core

At the moment, almost all LinOps in the public Ginkgo interface are templated only on ValueType and IndexType (except for preconditioner::Ilu/Ic where knowledge about the type of the trisolver is necessary, and the code is header-only), and any kind of composition between LinOps happens at the core level with separate kernels.
The batched setting would be rather different: Ideally, your preconditioner, SpMV and solver operations are merged into a small number of kernels on the reference/cuda/... kernel level. For this we likely need to explicitly instantiate a set of possible combinations. Now the question is, how do we do the dispatch? This problem has two layers:

1. dispatch for different batched formats:
   Based on the type of LinOp we pass to our BatchSolver (BatchDense, BatchCsr, ...), we need to select different kernels using the correct storage format. I think there is not much wiggle room here, if we don't want to introduce a fundamentally different way to apply batched operators.

2. dispatch for different solver and preconditioner types:
   Here we have a bit of freedom and can choose whether to use template parameters or factory parameters or something else.
   I would prefer for factory parameters for 4 reasons:
   2.1. we already need to set a few parameters like stopping criteria
   2.2. instantiating a new template for each type has more potential for binary code bloat
   2.3. having different types for different solver/preconditioner combinations makes the use of batch solvers as an abstract concept more complex, since we can't easily cast a batched LinOp to its concrete type without trying out a set of combinations or adding a common base class.
   2.4. using templates might significantly complicate the interaction between core and device code, since it is not clear to me right now how you pass on implementation information from SmallCsr to the concrete kernels for all executors.

Compile-time vs. run-time composition

If I understood correctly what you are describing in relation to unified storage, you wanted to have composition between SpMV, preconditioner and solver operations to happen at the run-time level with function pointers to __device__ functions, is that correct?
We should probably benchmark this at some point, but I would imagine that having indirect function calls on a platform without branch prediction might carry a certain performance penalty. Explicitly instantiating different combinations at compile-time would lead to more binary code, but might carry the added benefit of cross-function optimization, which nvcc especially is really agressive about.

On this note, I am not sure how useful a separate BatchPreconditioner concept would be, since this would only be used within a solver (if you really need just a simple preconditioner apply without kernel fusion, you could run a batched IR solver with the preconditioner and a single iteration only).

Batched vector vs. reusing Dense

About this one I am no longer entirely sure, since I see the following pros and cons pretty evenly balanced:

Batched vector:

• easier representation of batch-specific operations like batched axpy with different scalars, batched dot etc.
• batch boundaries also accessible without the batch operator or outside information from the application

Reusing Dense:

• Easier composability with other parts of Ginkgo that already use Dense as a vector type everywhere
• No need for batch size validation, which especially for the variable batch sizes can be non-negligible.

[Slaedr]

Thanks for putting all these considerations here. Indeed, I don't think templating BlockMatrix on the block storage format is strictly necessary for what I'm thinking. For example, an enum or a dynamic_cast could be used for functions to query the type of blocks at runtime, for the purpose of dispatching kernels, mostly. Note that at the moment I'm only thinking about uniform storage type for all the blocks in a BlockMatrix because I think that's the most useful case. Keeping the block types as runtime parameter(s) makes it easier to extend this if necessary in future.

Templates in Core

I agree with your points. But for enabling full optimization, at least the kernel would need to know the concrete type of the 'small' preconditioner, SpMV algorithm etc., as far as I can see.

Compile-time vs. run-time composition

Actually, it's the opposite. I was thinking of templating BlockPreconditioner on the 'small' solver type and the 'small' SpMV type so that these __device__ function calls can be inlined into the global preconditioner kernel. That's why I used the notation BlockPreconditioner<SmallSolver>. I don't really see a way around this if we want to avoid virtual function tables or branch statements in device code. That being said, branching may not be too bad in this case, because all threads in the thread-block will use the same preconditioner and therefore take the same branch. It might also be possible to simply 'hoist' the branch out of the kernel by the compiler, as happens in CPU code some times, in which case this branching is a complete non-issue.

Block preconditioners

Indeed, we don't really need a block preconditioner concept if we think purely in terms of batches. But if you generalize to block matrices and allow off-diagonal blocks, just preconditioning the diagonal blocks is not always enough. There needs to be a way to take off-diagonal blocks into account. This is where BlockPreconditioner<SmallSolver> comes in. In this scenario, our earlier conception of batch preconditioner corresponds to a block-Jacobi preconditioner. Note that 'block-Jacobi' here is not in the same sense as Ginkgo's block Jacobi. It simply refers to a global preconditioner that just approximately inverts the diagonal blocks. I have tried to adapt an earlier stub that I had written to illustrate this: https://www.overleaf.com/read/nzyvsqjdfknk

Batch vector vs. reusing dense

I don't have a strong opinion on this, but I slightly prefer to reuse dense, eg. this is how I wrote the Fbcsr class.

[yhmtsai]

One question: should LinOp need to support apply on Dense at least?
I think every dense/solver/preconditioner in the ginkgo current support on Dense.
I would like to make BatchLinOp to support on dense by default such that it can use solver or composition/combination of ginkgo.
And we still can add the BatchDense (or BatchVector) support to enable different batch-operation combination.
BatchDense -> Dense : I prefer to build the block-diagonal sense not concatenate the vectors to keep Batch in the same behavior. maybe add another conversion to allow concatenation.

[pratikvn]

I think there are two things here which are separate: Batching and Blocking.

• Blocking is a single LinOp, where data is blocked and the different blocks can and will interact with each other.
• Batching is execution of independent operations in parallel. Independent in the sense that there can be no interaction between the batches. This means it does not make any sense to have any off-diagonal blocks. Batching (augmentation) is a LinOp, but is just amalgamation of different individual LinOps which can be seen as one single LinOp for efficiency and performance reasons and is mathematically valid as well due to linearity.

In terms of interface, I think we should make a clear distinction between the two. Though from an implementation-al perspective, it makes sense to look at Batching as diagonal blocking, to me it makes no sense to implement Batching as a subset of Blocking.

One more thing to note is that you can have batches of blocked LinOps, which would become very convoluted if Batching were to be a subset of Blocking.

[tcojean]

Definitely agree here. There is indeed some non negligible overlap (in terms of software) between the two concepts and I actually at first thought we could unify the two like Aditya as well, but you make a good point that it's probably too restrictive to mix the two concepts and a potential source of issues since they actually have very different properties and potentially usage.

[Slaedr]

My reasons for this discussion are:

1. It just seems to me that the direction we are heading for the batching now almost completely looks like a block-diagonal solve with sparse blocks.
2. It would be nice to have generic block structure for the other applications that I talked about. Even the "CRI" problem, that created the requirement for batched solvers, really just seems like a special case. Presumably because of the way the time discretization is structured, the interaction between the cells appears to be treated explicitly. Therefore the reaction system is local to each cell. A different temporal discretization, where not only the chemistry but also the transport is treated implicitly, you would have off-diagonal blocks. So the main application that appears to demand development of batched solvers really seems like a special case of block solvers. The individual linear systems only happen to be independent - ultimately they're part of the same physical problem on the same mesh.
3. Suppose we do create a separate system of batch solvers now. Once we have a general block capability in future for maybe other applications like finite elements, we might have a lot of redundancy and code duplication. Unless there's a performance inefficiency incurred by regarding a batch of matrices as a literal block-diagonal matrix, it might be advantageous to do so from a software point of view. We then only need some helper functions to read a batch of matrices and convert it to a block-diagonal matrix.

[pratikvn]

I think in any case, if we have off-diagonal blocks, in the "CRI" case or otherwise, it would not batching, because then there would be dependencies between the batches. To me, by definition, batching means that there are no dependencies between the batches.

I agree that there are undoubtedly some commonalities between batching and blocking. It might even make sense to create some common base classes for both batching and blocking, but I don't think one can be a subset of the other.

Given that the blocking problem is a complex one and that we need the batching interface quite soon and in ginkgo 2.0 we could possibly overhaul the interface and create common base classes for blocking and batching (possibly with accessors, as you propose), I would suggest that we go ahead with a simple batching interface for now, without interleaving it with blocking and get back to blocking when 2.0 is nigh.