Pass to automatically reduce the number of new sparse tensors allocated

[paul-tqh-nguyen]

Most of our ops create new tensors as the output tensor via the empty_like function, the dup_tensor function, etc. We can also create new tensors via sparse_tensor.init op.

In some complicated use cases (e.g. GraphSAGE), I'm imagining it's possible that there are a lot of cases where we allocate a bunch of sparse tensors and delete them.

It would make sense that we re-use tensors when possible (e.g. once we're done with one and need to create one later, we'd just resize the old one instead of deleting it and creating a new one). This might lead to better memory efficiency, which might also lead to better runtime speed. It'd also reduce developer burden.

I think that there are a lot of parallel concepts from register allocation that we could adapt towards this problem.

• If our use case has several ops that return 2D f64 sparse tensors, we could think of this as a case where our machine has a limited number of 2D f64 sparse tensor "registers" and could use register allocation-inspired algorithms to minimize the number of new 2D f64 sparse tensors we'd have to allocate.
• We'd have to think of our machine as having a several sets of sparse tensor "registers". We'd have one set of sparse tensor "registers" for each (rank, dtype) sparse tensor type used in our use case.

Caveats:

• I'm not actually sure how much benefit we'd get from this.
• I'm not sure whether or not the main MLIR branch will come up with some better solution that would obviate this work.

I don't think this idea is limited to sparse tensors. I think an analogous idea could be used for applications primarily using dense tensors. Would this be useful for Numba in anyway?

[paul-tqh-nguyen]

I might've been overzealous when proposing this idea initially. I think the solution might be much simpler. We simply want to have a pass to minimize the number of sparse tensors created.

• If we create a new sparse tensor inside of a loop, we should hoist out the allocation of that tensor and simply re-use that same struct in the body of the loop. We should ideally be able to do repeated hoists with nested loops. It's unclear if we should hoist out of scf.if statements (it might be reasonable to always replace new sparse tensor allocations inside of an scf.if with one outside if we happen to have an extra unused one).
• If we have two sparse tensors with different value types but the value types have the same bit width, we should be able to cast one to the other and end up using just one sparse tensor.
• This might be a step too far, but if two sparse tensor have the same pointer and index bit widths but different value types, would it be reasonable to use one struct the whole time? For example, it seems possible to use an i64 tensor for a while and then do some casting to use it as an f32 tensor.
• Would it be reasonable for all of our ops to take an optional output tensor arg? This would reduce the number of sparse tensors initialized.

[paul-tqh-nguyen]

Here's an idea I had. I'm not sure if it's completely thought through, so I'm posting it here for a sanity check:

• Ideally, our memory-use optimization passes will be able to obviate the need for any user-written graphblas.update ops.
• Every place where a user writes a graphblas.update op, the user should ideally use graphblas.apply or graphblas.union or graphblas.intersect instead. Ideally, the optimization pass will transform these later ops into graphblas.update ops in order to have better memory use.