API: Add `permutedims(::SwizzleArray)`

[mtsokol]

Hi @willow-ahrens,

With Finch.jl 0.6.17 I managed to run first element-wise ops (.*) with broadcasting in a @compile lazy mode! Here I wanted to clarify the eager mode.

AFAIR The current approach is:

• If a Finch.jl function receives (at least one) LazyTensor input, then the output is also LazyTensor (lazy mode)
• If all inputs are SwizzleArray/Tensor then it's an eager mode. Finch.jl will wrap it in lazy(), call desired op, and call compute to get Tensor output.

For the latter I got some:

MethodError: no method matching permutedims(::SwizzleArray...)

I see that eager.jl operates on Tensor input but in finch-tensor we made a decision that everything is a SwizzleArray.

Should I make eager.jl also accept SwizzleArray inputs? (Check permutedims in this PR) Or would you like to approach it differently?

---

P.S. I added ndims(::LazyTensor{T, N}) as calling jl.ndims(obj) is more convenient than jl.ndims(jl.typeof(obj)) IMO

Also, I started wondering, how to implement size(::Finch.LazyTensor{T, N}).

[codecov[bot]]

Codecov Report

Attention: Patch coverage is 16.66667% with 5 lines in your changes are missing coverage. Please review.

Project coverage is 76.21%. Comparing base (e21aca8) to head (a49530c).

Files	Patch %	Lines
src/interface/lazy.jl	16.66%	5 Missing :warning:

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #470 +/- ## ========================================== - Coverage 76.25% 76.21% -0.05% ========================================== Files 92 92 Lines 8828 8834 +6 ========================================== + Hits 6732 6733 +1 - Misses 2096 2101 +5 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[willow-ahrens]

Should I make eager.jl also accept SwizzleArray inputs?

In short, yes. However, not all eager operations should be routed through the lazy(compute) interface. The lazy interface has a high overhead. For example, I think get/setindex has a more direct implementation that finch is currently using that we should stick with, because those ops are usually inexpensive and the overhead of scheduling them is probably way higher than just doing the op.

The following functions have direct implementations

• getindex/setindex because they are inexpensive
• copyto! and dropdefaults! because they are fundamental operations that are called by Tensor constructors
• isequal and == because I want a stable source of truth for these two

permutedims feels like it is similar to copyto!, and I'd like to keep it's direct eager implementation. However, SwizzleArray is itself also a form of laziness. Therefore, I propose the following functionalities:

1. permutedims(LazyTensor) -> LazyTensor
2. swizzle(LazyTensor) -> LazyTensor
3. lazy(arr::Swizzle{<:Tensor}) -> permutedims(lazy(arr.body), arr.dims)
4. permutedims(Tensor, dims) -> eager permutation to Tensor
5. permutedims(arr::Swizzle, dims) -> swizzle(arr, dims...)

Does this feel reasonable? I think we need to adjust the current codebase to handle 2 and 5.

[mtsokol]

Works for me!

One question: Does it mean that getindex/setindex only have eager versions? Can these operations be placed inside @compile decorated function?

[willow-ahrens]

Oh, in my review, I forgot that we should also add swizzle(arr::LazyTensor, dims...) = permutedims(arr, dims)

[willow-ahrens]

Also, I started wondering, how to implement size(::Finch.LazyTensor{T, N}).

LazyTensors currently only track extrude, i.e. whether the size is equal to one. I do think they could also track size, but I'm not sure yet whether it's a good idea. I would leave this alone for now because I think it's clearer to only keep track of the information that we are using to decide what computation is requested. However, I agree that it certainly would be very possible to track the size of the tensors along with their extruded dimensions. I just think that size(LazyTensor) should error out for now with a helpful error message that says to call compute on the lazy tensor.

One question: Does it mean that getindex/setindex only have eager versions? Can these operations be placed inside @compile decorated function?

In the near term, these operations should call compute and materialize.

In the future, we should add windowing operations to our lazy language to support getindex and setindex. Currently, this is not implemented. One way to implement these would be to understand them as multiplication by a permutation matrix or a similar selection mask with special properties that the compiler can analyze. let's make an issue for this and ask @kylebd99 as well for advice.

[kylebd99]

An indexing operator definitely makes sense to add at some point. It matches the last operator from relational algebra that's missing in our high level language, the selection operator. In the long run, it'll be important to have this because pushdown of selection operations would be pretty fundamental for getting good performance when people do use index slicing.

[willow-ahrens]

would it make sense to capture indexing with permutation matrices? I ask because that's a more minimal change

[willow-ahrens]

i.e.

x[3:5] = [0 0 1 0 0 0;
          0 0 0 1 0 0;
          0 0 0 0 1 0] * x

[kylebd99]

In my opinion, this approach obfuscates the relative simplicity of the operation and might make it harder to reason about performance. But, it is more in keeping with the style of finch at the moment, so I'm not sure.

[willow-ahrens]

Yeah, it could make things more complicated than they are worth. When we get to the point that indexing needs to become lazy, maybe it's worth sketching out what different notations would look like for transformation, optimization and analysis of the query notation. For the next 3 weeks I think that eager index is fine though!

[mtsokol]

@willow-ahrens I added missing swizzle and permutedims implementations from your list.

I also added getindex and size with more verbose error messages.