feat: inherit scalar indexing functionality from GPUArraysCore #268

avik-pal · 2024-11-12T16:54:12Z

needs some tests before merging

Example Usage

julia> using Reactant

julia> using GPUArraysCore

julia> x_ra = ConcreteRArray(rand(3, 4))
3×4 ConcreteRArray{Float64, 2}:
 0.166621  0.415209   0.23444   0.225489
 0.323775  0.201456   0.885111  0.625804
 0.22719   0.0906565  0.244437  0.98303

julia> x_ra[1]
0.1666208268454895

julia> GPUArraysCore.allowscalar(false)

julia> x_ra[1]
ERROR: Scalar indexing is disallowed.
Invocation of getindex(::ConcreteRArray, ::Vararg{Int, N}) resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] errorscalar(op::String)
   @ GPUArraysCore /mnt/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
 [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
   @ GPUArraysCore /mnt/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
 [4] assertscalar(op::String)
   @ GPUArraysCore /mnt/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
 [5] getindex(a::ConcreteRArray{Float64, 2}, args::Int64)
   @ Reactant /mnt/software/lux/Reactant.jl/src/ConcreteRArray.jl:175
 [6] top-level scope
   @ REPL[5]:1
 [7] top-level scope
   @ none:1

julia> @allowscalar x_ra[1]
0.1666208268454895

On CPU no error is ever thrown unless the user manually opts in for no-scalar indexing.

fixes #232

avik-pal · 2024-11-12T17:10:48Z

@mofeing can you check if this helps your case where you saw the scalar indexing warnings

mofeing · 2024-11-12T18:00:15Z

we confirm that this remove the infinite warnings we had in our code. Thanks @avik-pal!

I would approve the PR but it seems like this is breaking the tests?

avik-pal · 2024-11-12T18:01:44Z

The x86 ones are broken since we dont have the binaries in-place.

But I still need to add some tests before merging

wsmoses · 2024-11-12T18:03:58Z

src/Reactant.jl

@@ -110,12 +111,19 @@ function __init__()
 end

 function set_default_backend(backend::XLA.Client)
+    if backend === XLA.backends["cpu"]


so this won't quite work because we can end up with both cpu and gpu tensors

For XLA Buffer I can do a check with buffer on cpu, but I couldn't figure out how to do it for TracedRArray.

One solution is to set the local_task_storage to ScalarAllowed for CPU when entering the compile function

yeah traced we should always err for there, because that has its own problems of accidentally splitting up tensor ops into a bunch of scalars, regardless of backend impl

TracedRArray doesn't know about the backend accelerator, be CPU, GPU or TPU. actually, HLO dialects neither know about which backend are they gonna run in.
@wsmoses correct me if I'm wrong but that step is done later in XLA when compiling HLO to native executable

I mean even if they did [and you're right they don't] we should err for traced

I asked to remove them for CPU because it doesn't make much sense to raise a warning for CPU and they pollute a loooot the stdout.

i think removing for cpu concretearray is fine, but the problem is that it will equally pollute the IR we compile on traced, so we should still warn (or allowscalar)

Changed the behavior to (default CUDA behavior):

Allowed by warn in REPL

Disallowed with error in scripts

Can be locally allowed without warning using @allowscalar

should we reexport allowscalar?

wsmoses · 2024-11-12T18:09:29Z

src/ConcreteRArray.jl

-        )
-        getindex_warned[] = true
-    end
+


Can we have this call a function of our own which calls gpuarrays assertscalar if it's loaded

We can, but it is extremely lightweight:

julia> @time_imports using GPUArraysCore 0.2 ms Adapt 0.4 ms GPUArraysCore

eh okay I'm fine with this then

wsmoses · 2024-11-13T01:52:59Z

test/basic.jl

@@ -16,7 +14,7 @@ using InteractiveUtils

    a = Reactant.ConcreteRArray(x)

-    c_res = sum(a)
+    c_res = @allowscalar sum(a)


Okay ideally this shouldn't be required. I feel like a load of a concretenumber/tracednumber itself should be automatically allowscalar

These work fine if the backend is CPU but the default implementation of sum will just loop over the indices which fails the GPU ci

Wait really, shouldn’t it fall back to a reduce?

if not this is definitely a bug

Oh sorry this is for the concretearray not traced

yeah we should still eventually make this a reduce, but for another time

wsmoses reviewed Nov 12, 2024

View reviewed changes

avik-pal force-pushed the ap/scalar_indexing branch from 4b39384 to 0ae7cb3 Compare November 13, 2024 00:09

avik-pal added 4 commits November 12, 2024 19:29

feat: inherit scalar indexing functionality from GPUArraysCore

d74de20

chore: run formatter

45a6851

fix: always warn inside tracing unless opt-out

ce7fd6f

chore: reexport @allowscalar

371a21f

avik-pal force-pushed the ap/scalar_indexing branch from 389c6c9 to 371a21f Compare November 13, 2024 00:30

avik-pal added 2 commits November 12, 2024 19:36

feat: add isapprox for array types

3542a5f

fix: test fixes for scalar indexing

f54a37b

avik-pal force-pushed the ap/scalar_indexing branch from df60d2f to f54a37b Compare November 13, 2024 01:08

fix: allow scalar indexing in gather

08078f0

wsmoses reviewed Nov 13, 2024

View reviewed changes

wsmoses approved these changes Nov 13, 2024

View reviewed changes

wsmoses merged commit 5a60501 into main Nov 13, 2024
21 of 34 checks passed

wsmoses deleted the ap/scalar_indexing branch November 13, 2024 20:08

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: inherit scalar indexing functionality from GPUArraysCore #268

feat: inherit scalar indexing functionality from GPUArraysCore #268

avik-pal commented Nov 12, 2024 •

edited

Loading

avik-pal commented Nov 12, 2024

mofeing commented Nov 12, 2024

avik-pal commented Nov 12, 2024

wsmoses Nov 12, 2024

avik-pal Nov 12, 2024

wsmoses Nov 12, 2024

mofeing Nov 12, 2024 •

edited

Loading

wsmoses Nov 12, 2024

mofeing Nov 12, 2024 •

edited

Loading

wsmoses Nov 12, 2024

avik-pal Nov 13, 2024

mofeing Nov 13, 2024

avik-pal Nov 13, 2024

wsmoses Nov 12, 2024

avik-pal Nov 12, 2024

wsmoses Nov 13, 2024

wsmoses Nov 13, 2024

avik-pal Nov 13, 2024

wsmoses Nov 13, 2024

wsmoses Nov 13, 2024

feat: inherit scalar indexing functionality from GPUArraysCore #268

feat: inherit scalar indexing functionality from GPUArraysCore #268

Conversation

avik-pal commented Nov 12, 2024 • edited Loading

Example Usage

avik-pal commented Nov 12, 2024

mofeing commented Nov 12, 2024

avik-pal commented Nov 12, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mofeing Nov 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mofeing Nov 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

avik-pal commented Nov 12, 2024 •

edited

Loading

mofeing Nov 12, 2024 •

edited

Loading

mofeing Nov 12, 2024 •

edited

Loading