Jutho / TensorKit.jl

A Julia package for large-scale tensor computations, with a hint of category theory
MIT License
235 stars 41 forks source link

BigFloat issue #152

Open ebelnikola opened 2 months ago

ebelnikola commented 2 months ago

Hi!

First of all, thank you for this wonderful package. It is a pleasure to use it.

I have noticed that the @tensor macro fails to perform a contraction when tensors contain BigFloat numbers. Here is a minimal example:

using TensorKit
A_ok=Tensor(randn,ℝ^2⊗ℝ^2⊗ℝ^2)
@tensor res1[-1,-2]:=A_ok[1,2,-1]*A_ok[1,2,-2]; 

A_failure=Tensor(randn,BigFloat,ℝ^2⊗ℝ^2⊗ℝ^2)
@tensor res2[-1,-2]:=A_failure[1,2,-1]*A_failure[1,2,-2];

In this example, the first contraction goes well and the second throws:

ERROR: UndefRefError: access to undefined reference
Stacktrace:
  [1] getindex
    @ ./essentials.jl:13 [inlined]
  [2] getindex
    @ ~/.julia/packages/StridedViews/dcnHM/src/stridedview.jl:97 [inlined]
  [3] macro expansion
    @ ~/.julia/packages/Strided/Eoybo/src/mapreduce.jl:359 [inlined]
  [4] macro expansion
    @ ./simdloop.jl:77 [inlined]
  [5] macro expansion
    @ ~/.julia/packages/Strided/Eoybo/src/mapreduce.jl:358 [inlined]
  [6] _mapreduce_kernel!(f::Base.Fix2{…}, op::typeof(+), initop::Base.Fix2{…}, dims::Tuple{…}, blocks::Tuple{…}, arrays::Tuple{…}, strides::Tuple{…}, offsets::Tuple{…})
    @ Strided ~/.julia/packages/Strided/Eoybo/src/mapreduce.jl:229
  [7] _mapreduce_block!(f::Function, op::typeof(+), initop::Base.Fix2{…}, dims::Tuple{…}, strides::Tuple{…}, offsets::Tuple{…}, costs::Tuple{…}, arrays::Tuple{…})
    @ Strided ~/.julia/packages/Strided/Eoybo/src/mapreduce.jl:152
  [8] _mapreduce_order!(f::Function, op::Function, initop::Function, dims::Tuple{…}, strides::Tuple{…}, arrays::Tuple{…})
    @ Strided ~/.julia/packages/Strided/Eoybo/src/mapreduce.jl:138
  [9] _mapreduce_fuse!(f::Function, op::Function, initop::Function, dims::Tuple{…}, arrays::Tuple{…})
    @ Strided ~/.julia/packages/Strided/Eoybo/src/mapreduce.jl:116
 [10] _mapreducedim!(f::Function, op::Function, initop::Function, dims::Tuple{…}, arrays::Tuple{…})
    @ Strided ~/.julia/packages/Strided/Eoybo/src/mapreduce.jl:93
 [11] tensoradd!(C::StridedViews.StridedView{…}, pC::Tuple{…}, A::StridedViews.StridedView{…}, conjA::Symbol, α::Bool, β::Bool, backend::TensorOperations.Backend{…})
    @ TensorOperations ~/.julia/packages/TensorOperations/dNaBM/src/implementation/strided.jl:17
 [12] tensoradd!
    @ ~/.julia/packages/TensorOperations/dNaBM/src/implementation/strided.jl:8 [inlined]
 [13] _add_trivial_kernel!(::TrivialTensorMap{…}, ::TensorMap{…}, ::Tuple{…}, ::Function, ::Bool, ::Bool)
    @ TensorKit ~/.julia/packages/TensorKit/59D34/src/tensors/indexmanipulations.jl:334
 [14] add_transform!
    @ ~/.julia/packages/TensorKit/59D34/src/tensors/indexmanipulations.jl:323 [inlined]
 [15] add_permute!
    @ ~/.julia/packages/TensorKit/59D34/src/tensors/indexmanipulations.jl:276 [inlined]
 [16] permute!
    @ ~/.julia/packages/TensorKit/59D34/src/tensors/indexmanipulations.jl:16 [inlined]
 [17] permute(t::TensorMap{…}, ::Tuple{…}; copy::Bool)
    @ TensorKit ~/.julia/packages/TensorKit/59D34/src/tensors/indexmanipulations.jl:45
 [18] permute
    @ ~/.julia/packages/TensorKit/59D34/src/tensors/indexmanipulations.jl:31 [inlined]
 [19] _contract!(α::VectorInterface.One, A::TensorMap{…}, B::TensorMap{…}, β::VectorInterface.Zero, C::TensorMap{…}, oindA::Tuple{…}, cindA::Tuple{…}, oindB::Tuple{…}, cindB::Tuple{…}, p₁::Tuple{…}, p₂::Tuple{})
    @ TensorKit ~/.julia/packages/TensorKit/59D34/src/tensors/tensoroperations.jl:286
 [20] contract!(::TensorMap{…}, ::TensorMap{…}, ::Tuple{…}, ::TensorMap{…}, ::Tuple{…}, ::Tuple{…}, ::VectorInterface.One, ::VectorInterface.Zero)
    @ TensorKit ~/.julia/packages/TensorKit/59D34/src/tensors/tensoroperations.jl:254
 [21] tensorcontract!(::TensorMap{…}, ::Tuple{…}, ::TensorMap{…}, ::Tuple{…}, ::Symbol, ::TensorMap{…}, ::Tuple{…}, ::Symbol, ::VectorInterface.One, ::VectorInterface.Zero)
    @ TensorKit ~/.julia/packages/TensorKit/59D34/src/tensors/tensoroperations.jl:113
 [22] top-level scope
    @ ~/Codes/tmp/problem_example.jl:8

I also checked if the problem persists for plain arrays of BigFloat.

A_plain=randn(BigFloat,2,2,2);
@tensor res2[-1,-2]:=A_plain[1,2,-1]*A_plain[1,2,-2];

This works. I use Julia 1.10.5.

Update Here is what causes the problem. The function similar, when applied to tensors with BigFloat entries, gives a tensor with undefined entries. This leads to the following code failure with analogous error message:

using TensorKit

A=TensorMap(randn,BigFloat,ℝ^2←ℝ^2)
B=Tensor(randn,BigFloat,ℝ^2)
C=TensorKit.similar(A, promote_type(scalartype(A), scalartype(B)),TensorKit.compose(space(A),space(B)))

mul!(C,A,B)

It seems that _mapreduce_kernel! tries to use tensor elements of C for something.

lkdvos commented 2 months ago

You are right, the problem arises when the entries are not 'isbits' types, in which case they are initialized as undef. In these cases, the implementation of Strided indeed tries to access the data first, instead of just assigning, throwing an error. In TensorOperations I seem to recall we manually bypassed this by explicitly initialising the arrays of non-isbits types with zeros, presumably we can re-use that functionality here. I'll have a look tomorrow, thanks for bringing this up!

As a small side-note, there will probably be other things that are not fully compatible with BigFloat entries too. I don't think we have any LAPACK fallbacks, so factorisations will probably not work either, and that is less straightforward to fix.

lkdvos commented 2 months ago

I played around with some fixes (progress here), but it does not seem like it is too straightforward. I cannot say I have enough understanding of the inner workings of Strided to completely fix the problem (@Jutho might know more?), and I am not entirely sure this way of fixing it is ideal, as it requires explicitly initializing all of the BigFloat arrays with zeros. Additionally, we are currently in the process of rewriting a large part of the code-base, so it is not that easy to get a fix out to you quickly...

ebelnikola commented 2 months ago

I see, it is fine, there is no reason to hurry, thank you very much! The package works without problems with DoubleFloats.jl. For me, this is a more suitable number class anyway. As for the decompositions, I added the following workaround into the MatrixAlgebra module:

using GenericLinearAlgebra
function svd!(A::StridedMatrix{<:Number}, alg::Union{SVD,SDD}) 
    res = LinearAlgebra.svd!(A)
    return res.U, res.S, res.Vt
end

It seems to work, though I have not yet tested it very well. Could you please tell me if you see any immediate issues with this approach?

P.S. If at any moment support of non-isbits types will be important, I noticed that the same problem persists for addition and, I guess, for any operation that uses similar.

lkdvos commented 2 months ago

Thanks for also looking into it. Indeed, that looks like a good solution (which might automatically get incorporated in the near future, as a similar thing is required for CUDA anyways, see https://github.com/Jutho/TensorKit.jl/tree/ld-cuda). I would guess a similar solution is necessary/exists for QR, LQ, etc, for which you could take inspiration from there.

I am definitely interested in the extended precision things, and have not tried the GenericLinearAlgebra myself. If there are any more issues that pop up, feel free to let me know, I would like to keep this issue open and revisit this once I get the CUDA support and the new version up and running.