willow-ahrens / Finch.jl

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

Kbd-tensor-dot #465

Closed kylebd99 closed 5 months ago

kylebd99 commented 5 months ago

Per issue #464, this PR implements the tensordot function from the array interface.

To do this, I had to make some changes to the format inference code which I need a second pair of eyes on. Specifically, lines 189-223 of compute.jl.

The previous version of mapjoin format inference didn't take the indices of the arguments into account. For instance, in the matrix mult kernel, there's an interior mapjoin:

C[i,j,k] = A[i,j] * B[k,j]

The format inferred for C should have three levels, but it previously only had a two levels because it assumed both concordized and equal index sets (i.e. it assumed element-wise operations).

kylebd99 commented 5 months ago

Also, there was a small bug in the simplify_queries rewrite which is fixed on line 38 of compute.

kylebd99 commented 5 months ago

The tests which are failing right now seem like they demand slightly worse results than the current one (densifying more aggressively), but I feel like there might be a subtle reason why that behavior is necessary. For example, it expects the following,

 julia> 3 / A
Dense [:,1:3]
├─ [:, 1]: Dense [1:4]
│  ├─ [1]: Inf
│  ├─ [2]: 2.72727
│  ├─ [3]: 1.36364
│  └─ [4]: 0.909091
├─ [:, 2]: Dense [1:4]
│  ├─ [1]: Inf
│  ├─ [2]: Inf
│  ├─ [3]: Inf
│  └─ [4]: Inf
└─ [:, 3]: Dense [1:4]
   ├─ [1]: 0.681818
   ├─ [2]: Inf
   ├─ [3]: 0.545455
   └─ [4]: Inf

But this branch produces,


julia> 3 / A
Dense [:,1:3]
├─ [:, 1]: Sparse (Inf) [1:4]
│  ├─ [2]: 2.72727
│  ├─ [3]: 1.36364
│  └─ [4]: 0.909091
├─ [:, 2]: Sparse (Inf) [1:4]
└─ [:, 3]: Sparse (Inf) [1:4]
   ├─ [1]: 0.681818
   └─ [3]: 0.545455
willow-ahrens commented 5 months ago

@kylebd99 Okay, I think this is working! Regarding the test suite, I'll regenerate reference output and if anything looks like we can improve format inference, we'll leave it in! I'm gonna merge, and we can do the improved reordering heuristic in a separate PR. Thanks for your help!

codecov[bot] commented 5 months ago

Codecov Report

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

Project coverage is 76.16%. Comparing base (8f27be2) to head (75762e3).

Files Patch % Lines
src/interface/lazy.jl 85.71% 3 Missing :warning:
src/FinchLogic/nodes.jl 0.00% 2 Missing :warning:
src/interface/copy.jl 33.33% 2 Missing :warning:
src/interface/eager.jl 0.00% 2 Missing :warning:
Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #465 +/- ## ========================================== + Coverage 75.87% 76.16% +0.28% ========================================== Files 92 92 Lines 8800 8822 +22 ========================================== + Hits 6677 6719 +42 + Misses 2123 2103 -20 ```

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