Should batched operators be LinOps, or separate BatchLinOps?

[Slaedr]

Arguments in favour of separate BatchLinOp:

1. An apply signature like BatchLinOp* apply(const BatchLinOp*, BatchLinOp*) comes with a compile-time guard against trying to apply a batch operator to a non-batch operator etc. I think this is very useful for applications that use both batch functionality and regular functionality.
2. [Countered] If we don't have BatchLinOp, where would we put the validation of batch sizes of operators in apply? Note that the validation for batch operators is different from non-batch ones. We could keep it in EnableBatchLinOp, but we would still miss proper validation when LinOp::apply is used. My understanding is LinOp::apply is the one that gets called when one of the arguments is simply a LinOp* rather than a pointer to a concrete type.
   Answer using @upsj's idea: The batch validation could go into the backends and onto the device. This also lets us scalably do non-uniform batch operator validation.

Arguments against BatchLinOp (instead just keeping everything as LinOp):

1. Extra lines of code for class BatchLinOp and GKO_ENABLE_BATCH_LINOP_FACTORY, which repeat what is already done for the LinOp counterparts, are not necessary. (Note that we would still keep EnableBatchLinOp to keep batch_dim there.)
2. Having everything as a LinOp might make it easier to integrate task-based features later.

@pratikvn and @yhmtsai, could you document other points here?

[upsj]

Another point: High-Level Logging
Do we want to log BatchLinOp::apply operations (begin/end)? If so, we'd either need to add a new logger event, or reuse the LinOp one.

[pratikvn]

Yes, we currently add new events for BatchLinOp.

[pratikvn]

Copying stuff from our discussion on Slack:

I think there are a few differences both from a conceptual side and from an implementational side:

Conceptual:

1. In my opinion, LinOp is basically a single operator, while BatchLinOp is specifically a batch/collection of linear operators.
2. I think with point 1, it should be possible to see BatchLinOp as a parent of LinOp, but vice versa (even with the block-diagonal equivalence is a bit sketchy).
3. Currently, we of course just have all batches having the same linear operator, but I think technically this need not be the case. If say some matrices are harder to solve, it might make sense to have different preconditioners for some of them which would be more efficient than having complex preconditioners for all of the matrices in the batch.

Implementational:

1. One of the main reasons we went for BatchLinOp was the ability to have a separate batch_dim object and be able to store the different sizes and strides of the batches separately, even though currently we just use the same size for all matrices.
2. Deriving BatchLinOp from LinOp would complicate things. You would need to override the apply implementations which return a LinOp, so I think you would need dynamic_cast's everywhere ?
3. Another reason was that a Batched operators are experimental right now and if we tie them to the LinOp class hierarchy, and later decided that we would need a separate hierarchy, that would make it much harder than if after Batched operators are feature complete, we decided to move to the LinOp hierarchy, which would be much easier, in my opinion.

[upsj]

The discussion around concepts feels a bit arbitrary to me - BatchLinOp perfectly fits the description of a LinOp (i.e. it can be applied to other LinOps, e.g. a Dense vector), and every LinOp can be imagined as a BatchLinOp containing a single batch.
From a practical standpoint, LinOp is the more general type of the two, since BatchLinOp adds an additional member/functionality that LinOp does not have (batch count/batch sizes). This to me is an argument for LinOp being the base type. What issues would you expect from heterogeneous batches being based on LinOp?

On the user side, I would consider BatchLinOp::apply(Dense, Dense) an important use case, which isn't possible with the restriction to BatchLinOp parameters. We can have BatchVector specializations that store non-contiguous memory, and Dense for anything that stores its data in a packed, contiguous form.

[Slaedr]

About conceptual point 3, we don't actually have a way to do that and I don't see a straightforward way to do that even with BatchLinOp. Further, a "preconditioner" object with different preconditioners for different parts of the problem is still a preconditioner, and could still be thought of as a LinOp. Overall, I think I agree with @upsj that the conceptual arguments are somewhat arbitrary. A stack of independent matrices, for instance, could be interpreted as a 3-tensor, and a batch matrix vector product could be formulated as some kind of tensor multiplication, where the 3-tensor is now the linear operator in the usual sense. The caveats about Krylov methods not strictly being linear applies to both formulations, so that's immaterial here.

[Slaedr]

Implementation point 2: @yhmtsai I don't think that kind of code for BatchLinOp would work. We need to override LinOp::apply somehow to validate parameters in the batch sense. We would need something like EnableBatchLinOp that does that, and then it would need to eventually derive from LinOp. Still, my original issue in point 2 in favour of BatchLinOp (as it is right now in the PR) still holds.

[yhmtsai]

Yes, it will require all BatchLinOp apply to do their own validation on parameter because LinOp only check the entire size.
I am not sure scalar needs to be 1x1 or not

[Slaedr]

Actually from @upsj's comment on the PR, the validation may not be an issue at all. If we accept the discussion there, that removes this objection to everything being a LinOp.

[pratikvn]

After thinking about this a bit, I agree with @yhmtsai and I think it does make sense to go back to LinOp instead of a separate hierarchy for batched operators. Considering cases of tasked based approaches and distributed approaches in mind, given that LinOp is our usual base type, it would indeed make it a lot easier to have everything as LinOp.

I will update the PR, but before that I would also probably first test this on batch-develop to make sure we don't face any issues with our batch solvers and batch preconditioners.

[hartwiganzt]

Maybe it would be good if everyone involved in this discussion makes a list (on paper) with technical aspects motivating an independent BatchLinOp class and technical aspects motivating the use of LinOp. You may also include counterarguments if you have them. We then together create a superset of all arguments for both strategies. (In the best case, everybody already has this superset as he considered all aspects.) I understand this is some painful exercise, but maybe useful in this case?

[tcojean]

@yhmtsai But according to what you're saying, going a separate BatchLinOp does give some features: it prevents these cases you're talking about, things that cannot happen reasonably like a normal preconditioner with a batched solver (or almost anything with a batched solver, really) would explicitly not be possible because Batched wouldn't be a LinOp.

[yhmtsai]

I do not think about LinOp preconditioner into the BatchLinOp solver.
It should not happen. However, I think it should be the same cost to detect the preconditioner.
The BatchSolver needs to know the type of preconditioner.
Thus, no matter whether it is a BatchPrecond or LinopPrecond, they all have a list to search for that. LinOp Precond will not fit in the case.
For the LinOp Solver using BatchLinOp, it's useless in practice but should be perform correctly in theory.

[upsj]

@tcojean for Dense to be used with BatchLinOp, the interface would need to be BatchLinOp::apply(LinOp, LinOp) or require overloads for Dense and BatchLinOp.

Ginkgo has very little static type safety with regards to apply(...) and LinOpFactory, and we have no plans to change that so far, doing the same with BatchLinOp fits well into that framework. You can already pass a solver to a preconditioner factory, which leads to a runtime error that we throw explicitly in the LinOp/LinOpFactory implementation. As long as there are sensible LinOp - BatchLinOp interactions, I don't see a clean way to selectively allow only them.

From the block-diagonal viewpoint, BatchLinOp is a LinOp with a bit of additional information (block boundaries), just like DistributedLinOp has a communicator and local size, I am not sure whether it makes sense to treat them differently.

[tcojean]

The problem for me is from a programming POV, not mathematical or in the representation sense.

• Batched with Dense can easily be an overload or anything like that.
• When you use batched functionality, the user already has put all data on the device/wherever it should be, nothing should happen at all (no temporary clone, no host copy, etc) other than solving the problem in this exact setting.
• Because of the backend side, the batched functionality is very static in terms of what it supports, much more than our usual Ginkgo functionality. You can only combine its functionality with very specific settings, everything else would fail, and I believe that is the right way to go. The user wants a fast solution, and if the setting he asked for isn't available a clear error. Not to have conversions happen or other steps behind the scenes as we do everywhere in general Ginkgo's LinOp hierarchy so that he can use one of the normal functionality instead.
• I think putting this in LinOp isn't a clear message for the user gives a false impression that you can just combine all/most features whereas you either cannot or should not in almost all cases. In this way, most of Ginkgo is unneeded for the batched functionality. You only need the executor kernel dispatching support and hardware feature discovery, strong kernels, and a little bit of safety/feedback on top.

In general, I think it's a question of what type of interface we want, should it be more host side or have host options (LinOp), or should we target a more device-side interface? Arguably, we wouldn't even need BatchedLinOp for a good interface for this, something close to what Isha had done with simple base classes for solvers, preconditioners, matrices would be enough.

[Slaedr]

@tcojean On the relevance of LinOp vs BatchLinOp for tasking, I thought the tasking system might need some information on the task being submitted. For that one might add some (virtual) functions to LinOp so that the tasking system can query them (eg. some estimate of the memory bandwidth requirement as a function of num_rows and nnz). If BatchLinOp is a LinOp, that would be straightforward, and the task runtime itself does not really need to care whether some operation is a regular operation or batch operation. Or is that not how the tasking runtime would work? Would the tasking system rather be aware of batch operations and itself divide the batch into sub-batches?

Or perhaps you think the problem of getting the proper information to the task runtime is independent of what we do with BatchLinOps? Maybe we can have a Task feature which would provide the virtual functions the runtime needs, and any operation (whether LinOp or BatchLinOp) can inherit from that.

[tcojean]

@Slaedr The runtime in many cases because these are small batches with a fairly small amount of memory as well as operations, would not be able to recover the time it would take to move them or a few of them to another device. In that sense, you mostly don't need a runtime to begin with or would risk losing performance by switching from static scheduling (user defined) to dynamic (runtime controlled). If you had a really large amount of systems in the batch or if they were BLAS3 you might be able to do something though, but even that hasn't really been researched AFAIK.

From the runtime POV, it's entirely up in the air whether it would be able to reuse LinOp or anything else. The way I'm implementing the StarPU part currently, I only reuse the kernels of Ginkgo and need to rewrite the core algorithm for partitioning, to use task submission of kernels, play tricks on data (actual pointer/array) view, control the task submission flow, and whatnot. All the info you want then would need to be provided at the kernel level. For Batched if you would implement a task-based version, maybe it would instead exist on top of the BatchLinOp (or whatever) interface, so that you only play with "batches of batches".

[Slaedr]

So I think tasking can interact with batch solvers at two levels.

1. This is shared with other non-batched solvers. In presence of many different things for Ginkgo to do, like solve 2 different non-batched linear systems (eg. fluid momentum, fluid pressure) and one batched system (eg. chemical reactions in each cell), in addition to stuff the application does (eg. assemble each of those systems) and all this happens at every nonlinear step, the tasking system could orchestrate the whole thing, without interfering within any individual task.
2. The tasking system integrates with batch solvers. This is your point too. I'm thinking more like we submit the entire batch, and the tasking system divides it into sub-batches for optimal use of all GPUs on the system, and orchestrate that at the sub-batch level along with everything else going on in the application.

For point 1, it might be useful to have BatchLinOps as LinOps, so that that tasking system only needs to deal with the LinOp interface. Point 2 is more complicated, but is still consistent with having everything as LinOp, and EnableBatchLinOp or something providing just the batch interface.

[Slaedr]

Discussion:

BatchDense ("BatchVector") vs Dense ("Vector"):
- If we have BatchDense, we may not need BatchLinOp::apply(const Dense*,Dense*).
- Perhaps keep BatchDense for LHS matrices, but use simple Dense for vectors.

• Finally, because operations like dot and norm are different on "Vector" vs "BatchVector", the latter would be needed anyway.

Conclusion: There does not currently appear to be a need for interoperability between LinOp and BatchLinOp, and because having the separation makes application code more explicit and the development philosophy for batched functionality is a bit different, it is decided to keep BatchLinOp separate from LinOp. If at some point in the future, we realize we need an overall base class, we would break interface if necessary and make BatchLinOp a sub-type of LinOp.