willow-ahrens / Finch.jl

Sparse tensors in Julia and more! Datastructure-driven array programing language.
http://willowahrens.io/Finch.jl/
MIT License
151 stars 12 forks source link

Copying swizzles #609

Closed mtsokol closed 1 week ago

mtsokol commented 1 week ago

Hi @willow-ahrens,

I'm writing with one question regarding Tensor and SwizzleArray.

When I want to convert one format to another we have:

jl.swizzle(
    jl.Tensor(storage.levels_descr._obj, tensor._obj.body), *order
)

Which copies Tensor to the new format (and discards order of unwrapped tensor and sets order). This requires that order argument matches tensors's order.

I can also pass tensor's SwizzleArray that returns a Tensor with default F-major order:

jl.Tensor(storage.levels_descr._obj, tensor._obj)

I would like a new copy_swizzle(dest::SwizzleArray, src::SwizzleArray) function, which returns a SwizzleArray (or Tensor(dest::SwizzleArray, src::SwizzleArray) if it's more coherent).

This would allow me to pass CSC tensor and get a CSR array easily:

copy_swizzle(swizzle(Dense(SparseList(Element(0))), (2, 1)), csc_swizzle_array)

WDYT?

willow-ahrens commented 1 week ago

I think you can use copyto! here

mtsokol commented 1 week ago

@willow-ahrens Thank you! With copyto! resetting order is avoided. But I noticed that copyto! from dense to sparse format treats zeros as non-fill-values. How can I avoid it?

using Finch

arr = [0 1 0; 2 3 0]
tns = Tensor(arr)
sw = swizzle(tns, 1, 2)

copyto!(swizzle(Tensor(Dense(SparseList(Element(0)))), 1, 2), sw)

SwizzleArray (1, 2)
└─ 2×3-Tensor
   └─ Dense [:,1:3]
      ├─ [:, 1]: SparseList (0) [1:2]
      │  ├─ [1]: 0
      │  └─ [2]: 2
      ├─ [:, 2]: SparseList (0) [1:2]
      │  ├─ [1]: 1
      │  └─ [2]: 3
      └─ [:, 3]: SparseList (0) [1:2]
         ├─ [1]: 0
         └─ [2]: 0
mtsokol commented 1 week ago

I also managed to avoid copying fill-values with dropfills! but for csr format it fails with:

ERROR: Finch.FinchProtocolError("SparseListLevels cannot be updated multiple times")
using Finch

arr = [0 1 0; 2 3 0]
tns = Tensor(arr)
sw = swizzle(tns, 1, 2)

dropfills!(swizzle(Tensor(Dense(SparseList(Element(0)))), 2, 1), sw)
willow-ahrens commented 1 week ago

note: I have made a fix so that dropfills works like copyto. However, dropfills will always make a copy. You must still check for the cases where you can zero-copy during asarray.