Preserve fixed values in stage effect analysis

docs/src/learn/localmath-troubleshooting.md

@@ -52,6 +52,12 @@ unsafe globals, foreign calls, dynamic dispatch, recursion, and unsupported
 method shapes are rejected. The error reports the callable purpose, analyzed
 signature, selected method when available, and a recovery hint.
 
+Immutable fixed-shape storage values, including `StaticArray` vectors and
+matrices, remain scalar payloads during this analysis. Physical backend arrays
+are replaced by host surrogates. A `Reduce` operation must still infer its
+declared fixed value type, and every operation must satisfy the ordinary
+closed-effect contract.
+
 Real scalar `sin` and `cos` are available in ordinary evaluators. Their
 native floating-point methods use the same closed pure-call admission as
 `log` and `sqrt`; this is not permission for arbitrary foreign calls or

src/execution/stage_program_kernelabstractions.jl

@@ -752,7 +752,12 @@ function _ordering_effect_capability(
 end
 
 function _pointwise_surrogate_type(::Type{T}) where {T}
-    if T <: AbstractArray
+    if _storage_value_type(T)
+        # Immutable fixed-shape values are scalar payloads even though Julia's
+        # type hierarchy places StaticArrays below AbstractArray. Preserve the
+        # exact value signature while replacing only physical array storage.
+        return T
+    elseif T <: AbstractArray
         return Array{eltype(T), ndims(T)}
     elseif T <: NamedTuple
         names = T.parameters[1]

test/metal/stage_program.jl

@@ -2,6 +2,7 @@ using Test
 using Metal
 using LocalMath
 import KernelAbstractions
+import StaticArrays: SMatrix, SVector
 
 const LMSP = LocalMath
 
@@ -13,6 +14,12 @@ struct StageProgramParameterizedContribution end
 @inline (::StageProgramParameterizedContribution)(item::Int32, reads, parameters) =
     (value = LMSP.Contribution(item * getfield(parameters, 1)),)
 
+struct StageProgramFixedValueContribution{T}
+    value::T
+end
+@inline (evaluator::StageProgramFixedValueContribution)(item::Int32, reads, parameters) =
+    (value = LMSP.Contribution(evaluator.value),)
+
 struct StageProgramParameterizedCollect end
 @inline (::StageProgramParameterizedCollect)(item::Int32, reads, parameters) =
     (record = LMSP.CollectedValue(item * getfield(parameters, 1)),)
@@ -171,6 +178,36 @@ function stage_program_stage_program_reduce(backend)
         validated_generation, relation_status
 end
 
+function stage_program_fixed_value_reduce(backend)
+    return map((
+        SVector(1.0f0, -1.0f0),
+        SMatrix{2, 2}(1.0f0, -1.0f0, 2.0f0, -2.0f0),
+    )) do value
+        source = LMSP.Space(StageProgramStageProgramNode, 3)
+        destination = LMSP.Space(StageProgramStageProgramNode, 1)
+        output = LMSP.Field(destination, typeof(value))
+        relation = LMSP.FixedRelation(source => destination; degree = 1)
+        law = LMSP.Reduce(typeof(value), +;
+            seed = LMSP.IdentitySeed(zero(value)),
+            order = LMSP.CanonicalLeftFold())
+        publication = LMSP.Publication((LMSP.FieldPublication(
+            output, relation, LMSP.PublicationValue(:value)),), law)
+        stage = LMSP.Stage(source, NamedTuple(), (publication,),
+            LMSP.Evaluator(StageProgramFixedValueContribution(value)),
+            LMSP.Control(), LMSP.SourceOrigin(@__FILE__, @__LINE__;
+                label = :metal_fixed_value_reduce))
+        storage = Metal.MtlArray(fill(zero(value), 1))
+        endpoints = Metal.MtlArray(reshape(fill(Int32(1), 3), 1, 3))
+        counts = Metal.MtlArray(fill(Int32(1), 3))
+        bound = LMSP._bind_law(LMSP.LocalLaw(stage), LMSP._StructuralBinding(
+            (LMSP._field_storage_binding(output, storage),),
+            (LMSP._relation_storage_binding(relation, (; endpoints, counts)),)))
+        prepared = LMSP.prepare(LMSP.plan(bound; backend))
+        wait(LMSP.execute!(prepared))
+        return (actual = only(Array(storage)), expected = 3.0f0 * value)
+    end
+end
+
 function stage_program_stage_program_collect(backend)
     n = 513
     source = LMSP.Space(StageProgramStageProgramNode, n)
@@ -582,6 +619,13 @@ end
         @test LocalMath.inspect(prepared).stages[1].planning.executor === :collect
     end
 
+    selected in ("all", "fixed_value") &&
+    @testset "immutable fixed-value canonical Reduce" begin
+        for result in stage_program_fixed_value_reduce(backend)
+            @test result.actual == result.expected
+        end
+    end
+
     selected in ("all", "projection") &&
     @testset "candidate-owned Collect projection" begin
         projected, projected_count, positions =

test/test_reduce_stage.jl

@@ -1,5 +1,6 @@
 using Test
 import LocalMath
+import StaticArrays: SMatrix, SVector
 const LMR = LocalMath
 
 struct ReduceStageNode end
@@ -11,6 +12,11 @@ struct ItemContribution end
 struct FloatContribution end
 @inline (::FloatContribution)(item::Int32, reads, parameters) =
     (value = LMR.Contribution(Float32(item)),)
+struct FixedValueContribution{T}
+    value::T
+end
+@inline (evaluator::FixedValueContribution)(item::Int32, reads, parameters) =
+    (value = LMR.Contribution(evaluator.value),)
 struct TwoLaneContribution end
 @inline (::TwoLaneContribution)(item::Int32, reads, parameters) = (
     value = (LMR.Contribution(item),
@@ -78,6 +84,33 @@ end
     @test storage == [RecordContributionValue(UInt32(6), Int32(-6))]
 end
 
+@testset "canonical Reduce preserves immutable fixed values in exact typed IR" begin
+    for value in (
+            SVector(1.0f0, -1.0f0),
+            SMatrix{2, 2}(1.0f0, -1.0f0, 2.0f0, -2.0f0),
+        )
+        source = LMR.Space(ReduceStageNode, 3)
+        destination = LMR.Space(ReduceStageNode, 1)
+        value_type = typeof(value)
+        output = LMR.Field(destination, value_type)
+        relation = LMR.FixedRelation(source => destination; degree = 1)
+        law = LMR.Reduce(value_type, +;
+            seed = LMR.IdentitySeed(zero(value_type)),
+            order = LMR.CanonicalLeftFold())
+        stage = _reduce_test_stage(
+            source, output, relation, law,
+            FixedValueContribution(value))
+        storage = fill(zero(value_type), 1)
+        endpoints = reshape(fill(Int32(1), 3), 1, 3)
+        bound = LMR._bind_law(LMR.LocalLaw(stage), LMR._StructuralBinding(
+            (LMR._field_storage_binding(output, storage),),
+            (LMR._relation_storage_binding(relation, (
+                endpoints, counts = fill(Int32(1), 3))),)))
+        _run_test_candidate!(_prepare_test_candidate(bound))
+        @test storage == [3.0f0 * value]
+    end
+end
+
 @testset "canonical Reduce is the exact item-major lane-minor left fold" begin
     source = LMR.Space(ReduceStageNode, 3)
     destination = LMR.Space(ReduceStageNode, 2)