From 37838c6209cb6f0fbdc1486eb30c7d8316ee45bd Mon Sep 17 00:00:00 2001 From: PraneethMerugu Date: Wed, 9 Sep 2026 09:47:31 -0400 Subject: [PATCH 1/2] Preserve fixed values in stage effect analysis --- docs/src/learn/localmath-troubleshooting.md | 5 +++ .../stage_program_kernelabstractions.jl | 7 ++- test/metal/stage_program.jl | 44 +++++++++++++++++++ test/test_reduce_stage.jl | 33 ++++++++++++++ 4 files changed, 88 insertions(+), 1 deletion(-) diff --git a/docs/src/learn/localmath-troubleshooting.md b/docs/src/learn/localmath-troubleshooting.md index 7e6e32c..7913368 100644 --- a/docs/src/learn/localmath-troubleshooting.md +++ b/docs/src/learn/localmath-troubleshooting.md @@ -52,6 +52,11 @@ 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; an operation on a fixed value must still infer +that same fixed value type and 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 diff --git a/src/execution/stage_program_kernelabstractions.jl b/src/execution/stage_program_kernelabstractions.jl index fd9f274..7e7c56c 100644 --- a/src/execution/stage_program_kernelabstractions.jl +++ b/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] diff --git a/test/metal/stage_program.jl b/test/metal/stage_program.jl index e25fed5..5365055 100644 --- a/test/metal/stage_program.jl +++ b/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 = diff --git a/test/test_reduce_stage.jl b/test/test_reduce_stage.jl index 64fc1f4..8514c72 100644 --- a/test/test_reduce_stage.jl +++ b/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) From a1ddbca5c72cbc0291f11c3462bb8bfe2a1a9e61 Mon Sep 17 00:00:00 2001 From: PraneethMerugu Date: Wed, 9 Sep 2026 10:27:04 -0400 Subject: [PATCH 2/2] Clarify fixed-value reduction diagnostics --- docs/src/learn/localmath-troubleshooting.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/src/learn/localmath-troubleshooting.md b/docs/src/learn/localmath-troubleshooting.md index 7913368..3dbc689 100644 --- a/docs/src/learn/localmath-troubleshooting.md +++ b/docs/src/learn/localmath-troubleshooting.md @@ -54,8 +54,9 @@ 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; an operation on a fixed value must still infer -that same fixed value type and satisfy the ordinary closed-effect contract. +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