Fix scalar indexing on TileArrays and add codegen test (#94)

Author: AntonOresten

src/language/operations.jl

@@ -135,7 +135,7 @@ end
 @overlay function Base.getindex(arr::TileArray{T, N}, indices::Vararg{Integer, N}) where {T, N}
     tv = Intrinsics.make_tensor_view(arr)
     shape = ntuple(_ -> 1, Val(N))
-    pv = Intrinsics.make_partition_view(tv, Val(shape), PaddingMode.Undetermined)
+    pv = Intrinsics.make_partition_view(tv, shape, PaddingMode.Undetermined, nothing)
     tile = Intrinsics.load_partition_view(pv, nothing, true, promote(indices...) .- One())
     Intrinsics.to_scalar(reshape(tile, ()))
 end

test/codegen/operations.jl

@@ -1529,6 +1529,21 @@
             end
         end
 
+        @testset "TileArray scalar getindex" begin
+            @test @filecheck begin
+                @check_label "entry"
+                code_tiled(Tuple{ct.TileArray{Float32,1,spec1d}, ct.TileArray{Float32,1,spec1d}}) do lengths, out
+                    bid = ct.bid(1)
+                    @check "make_partition_view"
+                    @check "load_view_tko"
+                    @check "reshape"
+                    len = lengths[bid]
+                    ct.store(out, bid, ct.broadcast_to(ct.Tile(len), (16,)))
+                    return
+                end
+            end
+        end
+
         @testset "num_tiles helper" begin
             spec = ct.ArraySpec{2}(16, true)
             @test @filecheck begin

test/execution/basic.jl

@@ -1098,3 +1098,40 @@ end
 
     @test Array(b) ≈ Array(a)
 end
+
+@testset "scalar indexing as loop bound" begin
+    function scalar_index_loop_kernel(data::ct.TileArray{Float32,1},
+                                      lengths::ct.TileArray{Int32,1},
+                                      out::ct.TileArray{Float32,1})
+        bid = ct.bid(1)
+        len = lengths[bid]
+        acc = ct.zeros((16,), Float32)
+        j = Int32(1)
+        while j <= len
+            tile = ct.load(data, j, (16,))
+            acc = acc .+ tile
+            j += Int32(1)
+        end
+        ct.store(out, bid, acc)
+        return
+    end
+
+    # 3 blocks, each sums a different number of tiles
+    n_tiles = Int32[2, 3, 1]
+    data = CUDA.rand(Float32, 48)  # 3 tiles of 16
+    lengths = CuArray(n_tiles)
+    out = CUDA.zeros(Float32, 48)
+
+    ct.launch(scalar_index_loop_kernel, 3, data, lengths, out)
+
+    data_cpu = Array(data)
+    out_cpu = Array(out)
+    for bid in 1:3
+        expected = zeros(Float32, 16)
+        for j in 1:n_tiles[bid]
+            expected .+= data_cpu[(j-1)*16+1 : j*16]
+        end
+        @test out_cpu[(bid-1)*16+1 : bid*16] ≈ expected
+    end
+end
+