Add workaround for folding constant arithmetic.

Author: maleadt

src/language/operations.jl

@@ -111,21 +111,24 @@ end
 end
 
 """
-    store(arr::TileArray, index, tile::Tile) -> Nothing
+    store(arr::TileArray, index, tile::Tile) -> Tile
 
 Store a tile to a TileArray at the given index. Index is 1-indexed.
+Returns the stored tile (enables chaining and helps constant folding).
 """
 # Regular N-D tiles (N >= 1)
 @inline function store(arr::TileArray{T}, index, tile::Tile{T, Shape}) where {T, Shape}
     tv = Intrinsics.make_tensor_view(arr)
     pv = Intrinsics.make_partition_view(tv, Val(Shape), PaddingMode.Undetermined)
     Intrinsics.store_partition_view(pv, tile, (promote(index...) .- One())...)
+    return tile  # XXX: enables constant folding; remove when possible (see "constant folding" test)
 end
 
 @inline function store(arr::TileArray{T}, index::Integer, tile::Tile{T, Shape}) where {T, Shape}
     tv = Intrinsics.make_tensor_view(arr)
     pv = Intrinsics.make_partition_view(tv, Val(Shape), PaddingMode.Undetermined)
     Intrinsics.store_partition_view(pv, tile, index - One())
+    return tile  # XXX: enables constant folding; remove when possible (see "constant folding" test)
 end
 
 # Special case for 0D (scalar) tiles - reshape to 1D for partition view
@@ -135,13 +138,15 @@ end
     tile_1d = Intrinsics.reshape(tile, Val((1,)))
     pv = Intrinsics.make_partition_view(tv, Val((1,)), PaddingMode.Undetermined)
     Intrinsics.store_partition_view(pv, tile_1d, (promote(index...) .- One())...)
+    return tile  # XXX: enables constant folding; remove when possible (see "constant folding" test)
 end
 
 @inline function store(arr::TileArray{T}, index::Integer, tile::Tile{T, ()}) where {T}
     tv = Intrinsics.make_tensor_view(arr)
     tile_1d = Intrinsics.reshape(tile, Val((1,)))
     pv = Intrinsics.make_partition_view(tv, Val((1,)), PaddingMode.Undetermined)
     Intrinsics.store_partition_view(pv, tile_1d, index - One())
+    return tile  # XXX: enables constant folding; remove when possible (see "constant folding" test)
 end
 
 # Keyword argument version - dispatch to positional version

test/codegen.jl

@@ -1702,4 +1702,31 @@ end
             end
         end
     end
+
+    #=========================================================================
+     Constant Folding
+    =========================================================================#
+    @testset "constant folding" begin
+        spec = ct.ArraySpec{1}(16, true)
+
+        # XXX: This test verifies that store() returns the tile to enable constant
+        # folding. If this test fails after removing `return tile` from store(),
+        # Julia's optimizer will emit subi operations for constant index math.
+        # See operations.jl store() for the workaround.
+        @testset "store with constant index folds subtraction" begin
+            @test @filecheck begin
+                @check_label "entry"
+                @check "load_view_tko"
+                # Verify no subi appears between load and store - constant 1-1 should fold to 0
+                @check_not "subi"
+                @check "store_view_tko"
+                code_tiled(Tuple{ct.TileArray{Float32,1,spec}}) do a
+                    idx = Int32(1)
+                    tile = ct.load(a, idx, (16,))
+                    ct.store(a, idx, tile)
+                    return
+                end
+            end
+        end
+    end
 end