Keep tensor op descriptors out of constant AS

[maleadt]

Metal Shader Validation faults when __tensorops_impl_matmul2d_op_run receives a matmul descriptor initialized from a global that GPUCompiler cloned into AIR constant address space 2. GPUCompiler now provides a Metal hook for selecting the address space of each constant global, so override it for globals that feed tensor-op descriptor copies and leave only those constants in AS0 while preserving the default AS2 lowering for ordinary constants.

The inline tensor constructors also mark their explicitly packed-stride descriptors as non-contiguous, matching the initialization path used by the tensor-op runtime.

Fixes #834

[github-actions]

Metal Benchmarks

Details

Benchmark suite	Current: 5056aef	Previous: b702c8b	Ratio
array/accumulate/Float32/1d	821458 ns	821375 ns	1.00
array/accumulate/Float32/dims=1	1021687.5 ns	1023875 ns	1.00
array/accumulate/Float32/dims=1L	10535167 ns	10168458 ns	1.04
array/accumulate/Float32/dims=2	1302500 ns	1319750 ns	0.99
array/accumulate/Float32/dims=2L	6020208.5 ns	5141125 ns	1.17
array/accumulate/Int64/1d	1003375 ns	966729.5 ns	1.04
array/accumulate/Int64/dims=1	1152958 ns	1149208.5 ns	1.00
array/accumulate/Int64/dims=1L	12424937.5 ns	12139333 ns	1.02
array/accumulate/Int64/dims=2	1509354.5 ns	1499084 ns	1.01
array/accumulate/Int64/dims=2L	9341562 ns	9539583 ns	0.98
array/broadcast	355500 ns	385937.5 ns	0.92
array/construct	5667 ns	5666 ns	1.00
array/permutedims/2d	662209 ns	680333 ns	0.97
array/permutedims/3d	1141583.5 ns	1137208 ns	1.00
array/permutedims/4d	1570792 ns	1977958 ns	0.79
array/private/copy	420291.5 ns	445750 ns	0.94
array/private/copyto!/cpu_to_gpu	325833.5 ns	347188 ns	0.94
array/private/copyto!/gpu_to_cpu	322729 ns	382500 ns	0.84
array/private/copyto!/gpu_to_gpu	319125 ns	358667 ns	0.89
array/private/iteration/findall/bool	1162250 ns	1098792 ns	1.06
array/private/iteration/findall/int	1319229.5 ns	1265979.5 ns	1.04
array/private/iteration/findfirst/bool	1390500 ns	1479875 ns	0.94
array/private/iteration/findfirst/int	1447208 ns	1586959 ns	0.91
array/private/iteration/findmin/1d	1574583.5 ns	1608541.5 ns	0.98
array/private/iteration/findmin/2d	1297563 ns	1309542 ns	0.99
array/private/iteration/logical	1914396 ns	1776583 ns	1.08
array/private/iteration/scalar	1776708 ns	2849667 ns	0.62
array/random/rand/Float32	634750 ns	651854.5 ns	0.97
array/random/rand/Int64	688417 ns	684062.5 ns	1.01
array/random/rand!/Float32	540041 ns	591208 ns	0.91
array/random/rand!/Int64	496583 ns	522208 ns	0.95
array/random/randn/Float32	590584 ns	610417 ns	0.97
array/random/randn!/Float32	489500 ns	549708.5 ns	0.89
array/reductions/mapreduce/Float32/1d	502416 ns	781000 ns	0.64
array/reductions/mapreduce/Float32/dims=1	478833 ns	524209 ns	0.91
array/reductions/mapreduce/Float32/dims=1L	742229.5 ns	769042 ns	0.97
array/reductions/mapreduce/Float32/dims=2	482875 ns	535417 ns	0.90
array/reductions/mapreduce/Float32/dims=2L	1127520.5 ns	1361042 ns	0.83
array/reductions/mapreduce/Int64/1d	851416 ns	971875 ns	0.88
array/reductions/mapreduce/Int64/dims=1	780792 ns	813500 ns	0.96
array/reductions/mapreduce/Int64/dims=1L	1276250 ns	1515562.5 ns	0.84
array/reductions/mapreduce/Int64/dims=2	962958 ns	973167 ns	0.99
array/reductions/mapreduce/Int64/dims=2L	2235500 ns	2235000 ns	1.00
array/reductions/reduce/Float32/1d	497875 ns	773250 ns	0.64
array/reductions/reduce/Float32/dims=1	474792 ns	521833 ns	0.91
array/reductions/reduce/Float32/dims=1L	740625 ns	749437.5 ns	0.99
array/reductions/reduce/Float32/dims=2	477833 ns	533750 ns	0.90
array/reductions/reduce/Float32/dims=2L	1136166.5 ns	1356042 ns	0.84
array/reductions/reduce/Int64/1d	862500 ns	954667 ns	0.90
array/reductions/reduce/Int64/dims=1	785938 ns	813000 ns	0.97
array/reductions/reduce/Int64/dims=1L	1284875 ns	1467708 ns	0.88
array/reductions/reduce/Int64/dims=2	966854.5 ns	988604 ns	0.98
array/reductions/reduce/Int64/dims=2L	2218958 ns	2214125 ns	1.00
array/shared/copy	178125 ns	229292 ns	0.78
array/shared/copyto!/cpu_to_gpu	41833 ns	40542 ns	1.03
array/shared/copyto!/gpu_to_cpu	40625 ns	40417 ns	1.01
array/shared/copyto!/gpu_to_gpu	41542 ns	41209 ns	1.01
array/shared/iteration/findall/bool	1170208 ns	1105875 ns	1.06
array/shared/iteration/findall/int	1321000 ns	1268750 ns	1.04
array/shared/iteration/findfirst/bool	1146166 ns	1189270.5 ns	0.96
array/shared/iteration/findfirst/int	1183999.5 ns	1237333 ns	0.96
array/shared/iteration/findmin/1d	1335750 ns	1347625 ns	0.99
array/shared/iteration/findmin/2d	1305020.5 ns	1278437.5 ns	1.02
array/shared/iteration/logical	1771833.5 ns	1627625 ns	1.09
array/shared/iteration/scalar	5736 ns	5618 ns	1.02
integration/byval/reference	1185208 ns	1173542 ns	1.01
integration/byval/slices=1	1187167 ns	1168458 ns	1.02
integration/byval/slices=2	2131666 ns	2094042 ns	1.02
integration/byval/slices=3	18746125.5 ns	8211645.5 ns	2.28
integration/metaldevrt	472416 ns	465687.5 ns	1.01
kernel/indexing	348479 ns	384833 ns	0.91
kernel/indexing_checked	507833 ns	555291 ns	0.91
kernel/launch	13250 ns	13167 ns	1.01
kernel/rand	526667 ns	571458 ns	0.92
latency/import	1517691167 ns	1499339083.5 ns	1.01
latency/precompile	32382868166.5 ns	31947984208.5 ns	1.01
latency/ttfp	1836448375 ns	1811836708 ns	1.01
metal/synchronization/context	676.638036809816 ns	663.1736526946107 ns	1.02
metal/synchronization/stream	449.9191919191919 ns	434.67336683417085 ns	1.04

This comment was automatically generated by workflow using github-action-benchmark.

[maleadt]

Looks like JuliaGPU/GPUCompiler.jl#848 caused issues here.

[christiangnrd]

Nightly failure is JuliaGPU/GPUArrays.jl#735

[maleadt]

We triggered a different validation failure first, I think:

Error in testset examples/flashattention:
Error During Test at /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-macmini-aarch64-2.1/depots/5cd495a2-4a16-4674-ae02-c839447744bb/packages/ParallelTestRunner/kZFbj/src/ParallelTestRunner.jl:400
  Got exception outside of a @test
  LoadError: ArgumentError: Total used threadgroupMemoryLength(49152) must be <= 32768 bytes.

[christiangnrd]

The main :tensor tests weren't run but the flopscomp example (which ran first) tests the tensor (via :native) path as of my changes yesterday so this should be good. The flashattention failure I mention in #824 (comment) and I'm pretty sure it's the last validation error for CI to be green again