`LinearAlgebra.dot` and `x' * y` require `@allowscalar` but `sum(x .* y)` doesn't

[mofeing]

julia> a = rand(2)
2-element Vector{Float64}:
 0.41323466569127665
 0.4525169547907606

julia> b = rand(2)
2-element Vector{Float64}:
 0.37009493514138525
 0.818796709700481

julia> a_re = adapt(ConcreteRArray, a)
2-element ConcreteRArray{Float64, 1}:
 0.41323466569127665
 0.4525169547907606

julia> b_re = adapt(ConcreteRArray, b)
2-element ConcreteRArray{Float64, 1}:
 0.37009493514138525
 0.818796709700481

julia> f(x,y) = LinearAlgebra.dot(x,y)
f (generic function with 1 method)

julia> @jit f(a_re, b_re)
ERROR: Scalar indexing is disallowed.
Invocation of getindex(::TracedRArray, ::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:
...

julia> f(x,y) = x' * y
f (generic function with 1 method)

julia> @jit f(a_re, b_re)
ERROR: Scalar indexing is disallowed.
Invocation of getindex(::TracedRArray, ::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:
...

julia> f(x,y) = sum(x .* y)
f (generic function with 1 method)

julia> @jit f(a_re, b_re)
ConcreteRNumber{Float64}(0.5234554504635411)

[wsmoses]

I think we just need an overload for LinearAlgebra.dot. We have one for broadcast and sum so that makes sense

[AntonOresten]

LinearAlgebra.dot for vectors supported in #1213.

Seems x' * y works as well:

julia> Reactant.set_default_backend("gpu")

julia> a_re = Reactant.to_rarray(rand(2));

julia> b_re = Reactant.to_rarray(rand(2));

julia> f(x, y) = x' * y
f (generic function with 1 method)

julia> @jit f(a_re, b_re)
ConcretePJRTNumber{Float64, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(0.3989013436041066)

Worth mentioning dot works on any pair of AbstractArrays of the same length, and there's the three arg dot(x, A, y) version for dot(x, A*y).

[mofeing]

@AntonOresten do you mind checking if it works for 3 arg dot?

[AntonOresten]

It triggers scalar indexing:

julia> A_re = Reactant.to_rarray(rand(Float32, 2, 2))
julia> @jit dot(a_re, A_re, b_re)
ERROR: Scalar indexing is disallowed.
Invocation of getindex(::TracedRArray, ::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.

Any chance Reactant would fuse sum(x .* (A * y)) or something?

Regarding the AbstractArray case, it may be a simple case of flattening using vec and then reuse the vector method, since sum(x .* y) for arrays of different sizes could broadcast and inflate the result.

[AntonOresten]

Re 3 arg case:

Diagonal

julia> a = Reactant.to_rarray(rand(Float32, 100_000_000));

julia> b = Reactant.to_rarray(rand(Float32, 100_000_000));

julia> c = Reactant.to_rarray(Diagonal(rand(Float32, 100_000_000)));

julia> f(x) = sum(x)
f (generic function with 3 methods)

julia> f(x, y) = sum(x .* y)
f (generic function with 3 methods)

julia> f(x, A, y) = sum(x .* (A * y))
f (generic function with 3 methods)

julia> g = @compile f(a); @be Reactant.synchronize(g(a))
Benchmark: 107 samples with 1 evaluation
 min    743.437 μs (14 allocs: 400 bytes)
 median 759.503 μs (14 allocs: 400 bytes)
 mean   892.409 μs (14 allocs: 400 bytes)
 max    1.738 ms (14 allocs: 400 bytes)

julia> g = @compile f(a, b); @be Reactant.synchronize(g(a, b))
Benchmark: 65 samples with 1 evaluation
 min    1.339 ms (14 allocs: 400 bytes)
 median 1.446 ms (14 allocs: 400 bytes)
 mean   1.556 ms (14 allocs: 400 bytes)
 max    2.817 ms (14 allocs: 400 bytes)

julia> g = @compile f(a, c, b); @be Reactant.synchronize(g(a, c, b))
Benchmark: 50 samples with 1 evaluation
 min    1.934 ms (14 allocs: 400 bytes)
 median 2.001 ms (14 allocs: 400 bytes)
 mean   1.999 ms (14 allocs: 400 bytes)
 max    2.050 ms (14 allocs: 400 bytes)

julia> @code_hlo f(a, b)
module @reactant_f attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<100000000xf32>, %arg1: tensor<100000000xf32>) -> tensor<f32> {
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %0 = stablehlo.multiply %arg0, %arg1 : tensor<100000000xf32>
    %1 = stablehlo.reduce(%0 init: %cst) applies stablehlo.add across dimensions = [0] : (tensor<100000000xf32>, tensor<f32>) -> tensor<f32>
    return %1 : tensor<f32>
  }
}

julia> @code_hlo f(a, c, b)
module @reactant_f attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<100000000xf32>, %arg1: tensor<100000000xf32>, %arg2: tensor<100000000xf32>) -> tensor<f32> {
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %0 = stablehlo.multiply %arg1, %arg2 : tensor<100000000xf32>
    %1 = stablehlo.multiply %arg0, %0 : tensor<100000000xf32>
    %2 = stablehlo.reduce(%1 init: %cst) applies stablehlo.add across dimensions = [0] : (tensor<100000000xf32>, tensor<f32>) -> tensor<f32>
    return %2 : tensor<f32>
  }
}

Dense

Dense matmul will be dominating I guess so timing is not as interesting.

julia> a = Reactant.to_rarray(rand(Float32, 1000));

julia> b = Reactant.to_rarray(rand(Float32, 1000));

julia> c = Reactant.to_rarray(rand(Float32, 1000, 1000));

julia> @code_hlo f(a, c, b)
module @reactant_f attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<1000xf32>, %arg1: tensor<1000x1000xf32>, %arg2: tensor<1000xf32>) -> tensor<f32> {
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %0 = stablehlo.reshape %arg2 : (tensor<1000xf32>) -> tensor<1000x1xf32>
    %1 = stablehlo.dot_general %arg1, %0, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1000x1000xf32>, tensor<1000x1xf32>) -> tensor<1000x1xf32>
    %2 = stablehlo.reshape %1 : (tensor<1000x1xf32>) -> tensor<1000xf32>
    %3 = stablehlo.multiply %arg0, %2 : tensor<1000xf32>
    %4 = stablehlo.reduce(%3 init: %cst) applies stablehlo.add across dimensions = [0] : (tensor<1000xf32>, tensor<f32>) -> tensor<f32>
    return %4 : tensor<f32>
  }
}

These methods are a bit obscure, but wouldn't be surprised if they're used somewhere deep in someone's stack. LinearAlgebra defines a ton of methods for different pairings, so this is kind of a rabbit hole.