Support native trigonometric stages and checked canonical collection on Metal

CONTRIBUTING.md

@@ -97,6 +97,24 @@ Pointwise traversal and its control checks are owned by
 untouched backing storage, nonempty publication, and runtime-prefix diagnostics
 through the ordinary CPU and Metal inventories.
 
+Canonical collection ordering is owned by
+`src/execution/collect_physical_support.jl`. The shared
+`test/fixtures/collect_canonical_order_contracts.jl` checks reverse tuple-key
+ordering across workgroup boundaries, partial participation, empty gated
+collections, and duplicate-identity failure atomicity on CPU and real Metal.
+Run these checks with `--check-bounds=yes` when changing shared-memory sorting;
+disabled bounds checks must not substitute for correct padded-lane behavior.
+
+Closed callable admission is owned by
+`src/execution/stage_program_kernelabstractions.jl`. Its narrowly enumerated
+pure unary floating-point calls use public method reflection to require the same
+Base-owned native math module as the Float64 method, a concrete `AbstractFloat`
+argument and the same return type. Extensions undergo
+ordinary recursive effect analysis; this does not admit arbitrary foreign calls.
+`test/fixtures/trigonometric_stage_contracts.jl` exercises real sine/cosine
+publication through ordinary CPU and Metal stages, while
+`test/test_stage_preparation.jl` retains unsafe-capture/access rejection tests.
+
 Focused commands shorten the edit loop; they are not a second test inventory
 or release gate. Before handoff, run the complete suite of every changed
 package. Add the integration suite when a package boundary, extension, SciML

benchmark/collect_canonical_order.jl

@@ -0,0 +1,44 @@
+using Statistics
+import KernelAbstractions
+include("../test/fixtures/collect_canonical_order_contracts.jl")
+
+# Run in a normal package environment, or test/metal's environment with --metal.
+# Preparation and warmup are excluded; these timings are not pass/fail thresholds.
+backend = if "--metal" in ARGS
+    @eval import Metal
+    Metal.functional() || error("the selected Metal backend is not functional")
+    Metal.allowscalar(false)
+    Metal.MetalBackend()
+else
+    KernelAbstractions.CPU()
+end
+function benchmark_collect_order(backend, count)
+    (; prepared, storage) = _prepare_canonical_order_fixture(backend, count)
+    submit() = wait(
+        LocalMath.execute!(
+            prepared;
+            parameters = (; enabled = true, odd_only = false, duplicate = false, tied_keys = false)
+        )
+    )
+    for _ in 1:3
+        submit()
+    end
+    samples = map(1:9) do _
+        elapsed = @elapsed for _ in 1:10
+            submit()
+        end
+        elapsed / 10
+    end
+    records = collect(LocalMath.Adapt.adapt(Array, storage.records))
+    @assert Array(storage.count) == Int32[count]
+    @assert map(record -> record.slot, records) == Int32.(count:-1:1)
+    return println(
+        (;
+            count, minimum_seconds = minimum(samples),
+            median_seconds = median(samples), maximum_seconds = maximum(samples),
+        )
+    )
+end
+
+println("LocalMath=", pathof(LocalMath), " backend=", typeof(backend))
+foreach(count -> benchmark_collect_order(backend, count), (256, 4096))

docs/src/api/localmath.md

@@ -44,6 +44,13 @@ storage; and `allocate()` creates the exact bounded storage for a produced
 Collection. A caller-owned `StructArray` is borrowed unchanged. Allocating one
 copies its component arrays recursively and preserves the record layout.
 
+`Collect(...; order=canonical_by(key, identity))` orders participating records
+by their keys and identities, including tuple-valued keys. A closed participation
+gate publishes an empty logical collection without rewriting backing records.
+Duplicate canonical identities fail validation before changing the previously
+published count or records. The ordinary CPU and Metal collection-order tests
+exercise these behaviors across partial workgroups with bounds checks enabled.
+
 ## Public surface
 
 Ordinary authoring exports only the mathematical and execution vocabulary:

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.
 
+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
+side-effecting extensions of those functions. CPU/Metal comparisons use
+numerical tolerances rather than promising identical transcendental bits.
+
 ## Receipt or transaction failure
 
 `execute!` returns a logical `ExecutionReceipt`; call `wait` or `waitall` to

src/execution/collect_physical_support.jl

@@ -206,8 +206,11 @@ end
                 ascending = ((lane - 1) & sort_size) == 0
                 left = @inbounds local_order[lane]
                 right = @inbounds local_order[other]
-                swap = ascending ? _compacted_ordinal_less(port, workspace, right, left) :
-                    _compacted_ordinal_less(port, workspace, left, right)
+                # Keep one checked key-load call site in this shared-memory loop.
+                first_ordinal = ifelse(ascending, right, left)
+                second_ordinal = ifelse(ascending, left, right)
+                swap = _compacted_ordinal_less(port, workspace,
+                    first_ordinal, second_ordinal)
                 if swap
                     @inbounds begin
                         local_order[lane] = right

src/execution/stage_program_kernelabstractions.jl

@@ -295,7 +295,7 @@ const _POINTWISE_FAILURE_ONLY_INVOKES = (
     Base.error, Base.throw_boundserror, Core.throw_inexacterror,
 )
 const _POINTWISE_PURE_UNARY_FLOAT_INVOKES = (
-    Base.log, Base.sqrt, Base.cos,
+    Base.log, Base.sqrt, Base.sin, Base.cos,
 )
 const _POINTWISE_CALL_DEPTH_LIMIT = 32
 const _POINTWISE_AND_INT = getglobal(getglobal(Core, :Intrinsics), :and_int)
@@ -379,14 +379,27 @@ function _pointwise_residual_invoke_safe(
         return false
     end
     selected === method_instance.def || return false
-    if any(candidate -> candidate === binding,
-            _POINTWISE_PURE_UNARY_FLOAT_INVOKES)
-        length(signature.parameters) == 1 || return false
-        argument_type = only(signature.parameters)
-        argument_type isa DataType &&
-            argument_type <: AbstractFloat &&
-            isconcretetype(argument_type) || return false
-        return code_instance.rettype === argument_type
+    if any(
+            candidate -> candidate === binding,
+            _POINTWISE_PURE_UNARY_FLOAT_INVOKES
+        )
+        # Derive the native math owner through public method reflection. An
+        # extension replacing the Float64 reference must not grant its module
+        # the intrinsic shortcut for other floating-point types.
+        native_owner = try
+            parentmodule(binding, Tuple{Float64})
+        catch
+            nothing
+        end
+        if native_owner isa Module && parentmodule(native_owner) === Base &&
+                selected.module === native_owner
+            length(signature.parameters) == 1 || return false
+            argument_type = only(signature.parameters)
+            argument_type isa DataType &&
+                argument_type <: AbstractFloat &&
+                isconcretetype(argument_type) || return false
+            return code_instance.rettype === argument_type
+        end
     end
     lowered = try
         code_lowered(binding, signature)
@@ -500,6 +513,11 @@ function _pointwise_concrete_callable_invoke_safe(
     callable_type = first(spec.parameters)
     observed_type = _pointwise_operand_type(callee, analysis.typed_context)
     observed_type isa Type && observed_type <: callable_type || return false
+    if callee isa Function && Base.issingletontype(callable_type)
+        return _pointwise_residual_invoke_safe(
+            code_instance, callee, analysis, depth
+        )
+    end
     return _pointwise_method_instance_safe(method_instance, analysis, depth)
 end
 

test/fixtures/collect_canonical_order_contracts.jl

@@ -0,0 +1,103 @@
+using Test
+import LocalMath
+
+struct OrderedCollectionNode end
+struct OrderedCollectionRecord
+    slot::Int32
+    key::Tuple{Int32, UInt32, UInt32, UInt32}
+    identity::Tuple{UInt32, Int32, UInt32, Int32}
+end
+struct OrderedCollectionEvaluator end
+@inline function (::OrderedCollectionEvaluator)(item::Int32, reads, parameters)
+    odd_only, duplicate, tied_keys = parameters
+    slot = duplicate ? Int32(1) : item
+    record = OrderedCollectionRecord(
+        slot,
+        (tied_keys ? Int32(0) : -slot, UInt32(1), UInt32(0), UInt32(0)),
+        (UInt32(1), -slot, UInt32(0), Int32(1))
+    )
+    return (; records = LocalMath.CollectedValue(record, !odd_only || isodd(item)))
+end
+
+function _prepare_canonical_order_fixture(backend, count)
+    source = LocalMath.Space(OrderedCollectionNode, count)
+    collection = LocalMath.Collection(OrderedCollectionRecord, count)
+    enabled = LocalMath.Parameter(:enabled, Bool)
+    odd_only = LocalMath.Parameter(:odd_only, Bool)
+    duplicate = LocalMath.Parameter(:duplicate, Bool)
+    tied_keys = LocalMath.Parameter(:tied_keys, Bool)
+    publication = LocalMath.Publication(
+        (
+            LocalMath.CollectionPublication(
+                collection, LocalMath.PublicationValue(:records)
+            ),
+        ),
+        LocalMath.Collect(
+            OrderedCollectionRecord; maximum = 1,
+            order = LocalMath.canonical_by(:key, :identity)
+        )
+    )
+    stage = LocalMath.Stage(
+        source, NamedTuple(), (publication,),
+        LocalMath.Evaluator(OrderedCollectionEvaluator(), (odd_only, duplicate, tied_keys)),
+        LocalMath.Control(gate = enabled),
+        LocalMath.SourceOrigin(:canonical_collection, 1)
+    )
+    law = LocalMath.LocalLaw(
+        stage;
+        parameters = LocalMath.ParameterSchema(odd_only, duplicate, tied_keys, enabled)
+    )
+    storage = LocalMath.CompactedStorage(
+        backend, OrderedCollectionRecord, count;
+        source_items = count
+    )
+    prepared = LocalMath.prepare(law, collection => storage; backend)
+    return (; prepared, storage)
+end
+
+function _collect_canonical_order_contract(backend)
+    host_records(storage) = collect(LocalMath.Adapt.adapt(Array, storage.records))
+    return @testset "canonical tuple-key collection" begin
+        for count in (1, 2, 255, 256, 257)
+            @testset "capacity $count" begin
+                (; prepared, storage) = _prepare_canonical_order_fixture(backend, count)
+                @test Array(storage.count) == Int32[0]
+                run(; enabled = true, odd_only = false, duplicate = false, tied_keys = false) =
+                    wait(LocalMath.execute!(prepared; parameters = (; enabled, odd_only, duplicate, tied_keys)))
+                run()
+                @test Array(storage.count) == Int32[count]
+                @test map(record -> record.slot, host_records(storage)) == Int32.(count:-1:1)
+                run(; odd_only = true)
+                expected = reverse(filter(isodd, Int32.(1:count)))
+                @test Array(storage.count) == Int32[length(expected)]
+                @test map(
+                    record -> record.slot,
+                    host_records(storage)[1:length(expected)]
+                ) == expected
+                run(; tied_keys = true)
+                @test Array(storage.count) == Int32[count]
+                @test map(record -> record.slot, host_records(storage)) == Int32.(count:-1:1)
+                published = host_records(storage)
+                run(; enabled = false)
+                @test Array(storage.count) == Int32[0]
+                @test host_records(storage) == published
+                if count > 1
+                    run()
+                    published = host_records(storage)
+                    previous_count = Array(storage.count)
+                    failure = try
+                        run(; duplicate = true)
+                        nothing
+                    catch error
+                        error
+                    end
+                    @test failure isa LocalMath.LocalMathValidationError
+                    @test failure.contract === :runtime_stage_validation
+                    @test failure.actual.failure_class === :duplicate_identity
+                    @test Array(storage.count) == previous_count
+                    @test host_records(storage) == published
+                end
+            end
+        end
+    end
+end

test/fixtures/trigonometric_stage_contracts.jl

@@ -0,0 +1,51 @@
+using Test
+import LocalMath
+import KernelAbstractions
+
+struct TrigonometricStageDomain end
+struct TrigonometricStageEvaluator{F}
+    operation::F
+end
+@inline function (evaluator::TrigonometricStageEvaluator)(item::Int32, reads, parameters)
+    angle = something(reads[1][1].value)
+    return (value = LocalMath.UniqueValue(evaluator.operation(angle)),)
+end
+
+function trigonometric_stage_contracts(array_type, ::Type{T} = Float32) where {T}
+    return @testset "real trigonometric stage publication ($T)" begin
+        angles = T[-100, -3, -0.25, -0.0, 0, 0.25, 1, 3, 100]
+        space = LocalMath.Space(TrigonometricStageDomain, length(angles))
+        input = LocalMath.Field(space, T)
+        output = LocalMath.Field(space, T)
+        relation = LocalMath.IdentityRelation(space)
+        for operation in (sin, cos)
+            stage = LocalMath.Stage(
+                space, (angle = LocalMath.Access(input, relation),),
+                (
+                    LocalMath.Publication(
+                        (LocalMath.FieldPublication(output, relation, LocalMath.PublicationValue(:value)),),
+                        LocalMath.Unique(T),
+                    ),
+                ),
+                LocalMath.Evaluator(TrigonometricStageEvaluator(operation), ()),
+                LocalMath.Control(),
+                LocalMath.SourceOrigin(@__FILE__, @__LINE__; label = :trigonometric_publication),
+            )
+            values = array_type(angles)
+            destination = array_type(fill(T(17), length(angles)))
+            prepared = LocalMath.prepare(
+                LocalMath.LocalLaw(stage), input => values, output => destination;
+                backend = KernelAbstractions.get_backend(values),
+            )
+            @test Array(destination) == fill(T(17), length(angles))
+            wait(LocalMath.execute!(prepared))
+            # Higher-precision host arithmetic is a numerical oracle, not a
+            # second stage evaluator or a device bitwise-transcendental claim.
+            reference = setprecision(BigFloat, 128) do
+                T.(operation.(BigFloat.(angles)))
+            end
+            @test Array(destination) ≈ reference rtol = 4eps(T) atol = 4eps(T)
+            @test isequal(Array(values), angles)
+        end
+    end
+end

test/metal/collect_canonical_order.jl

@@ -0,0 +1,7 @@
+using Test
+import Metal
+include("../fixtures/collect_canonical_order_contracts.jl")
+
+Metal.functional() || error("canonical collection checks require real Metal")
+Metal.allowscalar(false)
+_collect_canonical_order_contract(Metal.MetalBackend())

test/metal/runtests.jl

@@ -16,6 +16,8 @@ const LOCALMATH_METAL_WITNESSES = (
     "product_values.jl",
     "ordered_fold_control.jl",
     "empty_pointwise_domains.jl",
+    "collect_canonical_order.jl",
+    "trigonometric_stages.jl",
 )
 
 @testset "LocalMath Metal runner inventory" begin

test/metal/trigonometric_stages.jl

@@ -0,0 +1,2 @@
+include(joinpath(@__DIR__, "..", "fixtures", "trigonometric_stage_contracts.jl"))
+trigonometric_stage_contracts(Metal.MtlArray, Float32)

test/runtests.jl

@@ -11,6 +11,7 @@ const LOCALMATH_INCLUDED_TESTS = (
     "test_stage_preparation.jl",
     "test_direct_pointwise_stage.jl",
     "test_empty_pointwise_domains.jl",
+    "test_trigonometric_stages.jl",
     "test_product_values.jl",
     "test_unique_stage.jl",
     "test_stage_program_lifecycle.jl",
@@ -21,6 +22,7 @@ const LOCALMATH_INCLUDED_TESTS = (
     "test_candidate_grouping.jl",
     "test_collect_stage_model.jl",
     "test_collect_stage_execution.jl",
+    "test_collect_canonical_order.jl",
     "test_ordered_fold_stage_model.jl",
     "test_ordered_fold_stage_execution.jl",
     "test_ordered_fold_control.jl",

test/test_collect_canonical_order.jl

@@ -0,0 +1,5 @@
+using Test
+import KernelAbstractions
+include("fixtures/collect_canonical_order_contracts.jl")
+
+_collect_canonical_order_contract(KernelAbstractions.CPU())