Tullio.jl Scalar indexing error when using Metal GPU backend

I've been porting a simulation to use the GPU. There is just one piece that I'm missing to keep everyting on the GPU side, a reduction using Tullio

# For dimensions >= 2
@inline function reduce_forces!(Fij, Dijd, Pij, ::Val{D}) where {D}
    @tullio Fij[d, i] = Dijd[i, j, d] * Pij[i, j]
    return nothing
end

My local machine can't use CUDA, so I'm testing with Metal.jl and I haven't been able to get Tullio to work with Metal arrays. Here is a MWE

using Tullio, KernelAbstractions, Metal
mul(A, B, C) = @tullio C[k] = A[k] * B[k];  # run macro after loading packages

n = 2^9
A = rand(n, n)
B = rand(n, n)
C = similar(A)

# Works on CPU
mul(A, B, C)

# Throws "scalar indexing" error on Metal backend, not on CUDA though
mul(mtl(A), mtl(B), mtl(C))

Here's the full error:

julia> mul(mtl(A), mtl(B), mtl(C))
ERROR: Scalar indexing is disallowed.
Invocation of getindex 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 ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
  [4] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
  [5] getindex
    @ ~/.julia/packages/GPUArrays/ZRk7Q/src/host/indexing.jl:50 [inlined]
  [6] 𝒜𝒸𝓉!
    @ ~/.julia/packages/Tullio/2zyFP/src/macro.jl:1036 [inlined]
  [7] tile_halves(fun!::var"#𝒜𝒸𝓉!#1", ::Type{…}, As::Tuple{…}, Is::Tuple{…}, Js::Tuple{}, keep::Nothing, final::Bool)
    @ Tullio ~/.julia/packages/Tullio/2zyFP/src/threads.jl:139
  [8] tile_halves(fun!::var"#𝒜𝒸𝓉!#1", ::Type{…}, As::Tuple{…}, Is::Tuple{…}, Js::Tuple{}, keep::Nothing, final::Bool) (repeats 12 times)
    @ Tullio ~/.julia/packages/Tullio/2zyFP/src/threads.jl:142
  [9] tile_halves
    @ ~/.julia/packages/Tullio/2zyFP/src/threads.jl:136 [inlined]
 [10] threader
    @ ~/.julia/packages/Tullio/2zyFP/src/threads.jl:65 [inlined]
 [11] macro expansion
    @ ~/.julia/packages/Tullio/2zyFP/src/macro.jl:1004 [inlined]
 [12] mul(A::MtlMatrix{…}, B::MtlMatrix{…}, C::MtlMatrix{…})
    @ Main ./REPL[2]:1
 [13] top-level scope
    @ REPL[8]:1
 [14] top-level scope
    @ ~/.julia/packages/Metal/UahNC/src/initialization.jl:80
Some type information was truncated. Use `show(err)` to see complete types.

Thanks in advance for the help, and thanks for the work on the package in general

Oct 03 '25 12:10 alonsoC1s

I can reproduce this, but don't know why it's failing.

The next thing to check is probably whether the same error occurs with KernelAbstractions alone. Expanding the macro, or reading what it says with verbose=2, it should be possible to reconstruct the functions. This readme section has some explanation of what act! is meant to do.

julia> mul(mtl(A), mtl(B), mtl(C))
ERROR: Scalar indexing is disallowed.
...
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] errorscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
...
  [6] 𝒜𝒸𝓉!
    @ ~/.julia/packages/Tullio/2zyFP/src/macro.jl:1036 [inlined]
  [7] tile_halves(fun!::var"#𝒜𝒸𝓉!#5", ::Type{MtlMatrix{Float32, Metal.PrivateStorage}}, As::Tuple{MtlMatrix{Float32, Metal.PrivateStorage}, MtlMatrix{Float32, Metal.PrivateStorage}, MtlMatrix{Float32, Metal.PrivateStorage}}, Is::Tuple{UnitRange{Int64}}, Js::Tuple{}, keep::Nothing, final::Bool)
    @ Tullio ~/.julia/packages/Tullio/2zyFP/src/threads.jl:139
...
 [12] mul(A::MtlMatrix{…}, B::MtlMatrix{…}, C::MtlMatrix{…})
    @ Main ./REPL[77]:1
 [13] top-level scope
    @ REPL[91]:1
 [14] top-level scope
    @ ~/.julia/packages/Metal/UahNC/src/initialization.jl:80
Some type information was truncated. Use `show(err)` to see complete types.

(jl_SJaIcd) pkg> st
Status `/private/var/folders/yq/4p2zwd614y59gszh7y9ypyhh0000gn/T/jl_SJaIcd/Project.toml`
  [63c18a36] KernelAbstractions v0.9.38
  [dde4c033] Metal v1.8.1
  [bc48ee85] Tullio v0.3.8

julia> @tullio C[k] = A[k] * B[k] verbose=2 grad=false;
┌ Info: >>>>> Preliminary expressions
│   verbosetidy(ex_pre) =
│    quote
└    end
┌ Info: ===== Base actor 
│   verbosetidy(ex_act) =
│    quote
│        local @inline(function 𝒜𝒸𝓉!(::Type, ℛ::AbstractArray{𝒯}, A, B, 𝒶𝓍k, ♻️ = nothing, 💀 = true) where 𝒯
│                    @inbounds @fastmath(begin
│                                for k = 𝒶𝓍k
│                                    ℛ[k] = A[k] * B[k]
│                                end
│                            end)
│                end)
└    end
┌ Info: =====KA===== KernelAbstractions kernel 
│   verbosetidy(kex1) =
│    quote
│        KernelAbstractions.@kernel function var"##🇨🇺#248"(ℛ::AbstractArray{𝒯}, A, B, 𝒶𝓍k, ♻️, 💀) where 𝒯
│                (k,) = @index(Global, NTuple)
│                begin
│                    ℛ[k] = A[k] * B[k]
│                end
│            end
└    end
[ Info: success expanding KernelAbstractions.@kernel
┌ Info: threading threshold (from cost = 1)
└   block = 262144
┌ Info: In-place body
│   verbosetidy(ex_body) =
│    quote
│        local 𝒶𝓍k = (Tullio.linearindex)(C)
│        (Tullio.linearindex)(A) == (Tullio.linearindex)(C) || (throw)("range of index k must agree")
│        (Tullio.linearindex)(B) == (Tullio.linearindex)(C) || (throw)("range of index k must agree")
│        @info "left index ranges" k = (Tullio.linearindex)(C)
│        begin
│            (Tullio.threader)(𝒜𝒸𝓉!, (Tullio.storage_type)(C, A, B), C, tuple(A, B), tuple(𝒶𝓍k), tuple(), +, 262144, nothing)
│            C
│        end
└    end

Oct 03 '25 21:10 mcabbott

Ah, a possible culprit is this:

https://github.com/mcabbott/Tullio.jl/blob/master/ext/TullioCUDAExt.jl

The CPU path has recursive tiling before calling act!, but the GPU path is meant to skip that. But tile_halves in stacktrace above hints it's not. And perhaps the problem is that this is done only for T<:CUDA.CuArray.

It's possible that there are other such mistakes. In fact it's possible that nothing ever worked with Metal.jl.

Sorry I don't have much time to fix things around here. But a PR making it do for Metal what it does for CUDA might not be too hard.

Oct 03 '25 21:10 mcabbott

I'd love to take a closer look, I've been meaning to contribute for a while I'll recreate the extension with MtlArray and see what happens

Oct 04 '25 08:10 alonsoC1s