Metal: add sincos parameter attributes (#762)

Co-authored-by: Tim Besard <tim.besard@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: 3 people

src/metal.jl

@@ -1706,50 +1706,68 @@ function annotate_air_intrinsics!(@nospecialize(job::CompilerJob), mod::LLVM.Mod
         isdeclaration(f) || continue
         fn = LLVM.name(f)
 
-        attrs = function_attributes(f)
-        function add_attributes(names...)
+        fn_attrs = function_attributes(f)
+        function add_fn_attributes(names...)
             for name in names
                 if LLVM.version() >= v"16" && name in ["argmemonly", "inaccessiblememonly",
                                                        "inaccessiblemem_or_argmemonly",
                                                        "readnone", "readonly", "writeonly"]
                     # XXX: workaround for changes from https://reviews.llvm.org/D135780
                     continue
                 end
-                push!(attrs, EnumAttribute(name, 0))
+                push!(fn_attrs, EnumAttribute(name, 0))
+            end
+            changed = true
+        end
+
+        function add_param_attributes(idx, names...)
+            param_attrs = parameter_attributes(f, idx)
+            for name in names
+                if name == "nocapture" && LLVM.version() >= v"21"
+                    # `nocapture` was replaced by `captures(none)` in LLVM 21 (an
+                    # integer-valued IntAttr, value 0 == CaptureInfo::none()).
+                    push!(param_attrs, EnumAttribute("captures", 0))
+                else
+                    push!(param_attrs, EnumAttribute(name, 0))
+                end
             end
             changed = true
         end
 
         # synchronization
         if fn == "air.wg.barrier" || fn == "air.simdgroup.barrier"
-            add_attributes("nounwind", "mustprogress", "convergent", "willreturn")
+            add_fn_attributes("nounwind", "mustprogress", "convergent", "willreturn")
+
+        # sincos
+        elseif match(r"^air.(fast_)?sincos", fn) !== nothing
+            add_param_attributes(2, "nocapture", "writeonly")
 
         # atomics
         elseif match(r"air.atomic.(local|global).load", fn) !== nothing
             # TODO: "memory(argmem: read)" on LLVM 16+
-            add_attributes("argmemonly", "readonly", "nounwind")
+            add_fn_attributes("argmemonly", "readonly", "nounwind")
         elseif match(r"air.atomic.(local|global).store", fn) !== nothing
             # TODO: "memory(argmem: write)" on LLVM 16+
-            add_attributes("argmemonly", "writeonly", "nounwind")
+            add_fn_attributes("argmemonly", "writeonly", "nounwind")
         elseif match(r"air.atomic.(local|global).(xchg|cmpxchg)", fn) !== nothing
             # TODO: "memory(argmem: readwrite)" on LLVM 16+
-            add_attributes("argmemonly", "nounwind")
+            add_fn_attributes("argmemonly", "nounwind")
         elseif match(r"^air.atomic.(local|global).(add|sub|min|max|and|or|xor)", fn) !== nothing
             # TODO: "memory(argmem: readwrite)" on LLVM 16+
-            add_attributes("argmemonly", "nounwind")
+            add_fn_attributes("argmemonly", "nounwind")
 
         # simdgroup
         elseif match(r"air.simdgroup_matrix_8x8_multiply_accumulate", fn) !== nothing
-            add_attributes("convergent", "mustprogress", "nounwind", "willreturn")
+            add_fn_attributes("convergent", "mustprogress", "nounwind", "willreturn")
         elseif match(r"air.simdgroup_matrix_8x8_load", fn) !== nothing
-            add_attributes("convergent", "mustprogress", "nofree", "nounwind", "readonly", "willreturn")
+            add_fn_attributes("convergent", "mustprogress", "nofree", "nounwind", "readonly", "willreturn")
         elseif match(r"air.simdgroup_matrix_8x8_store", fn) !== nothing
-            add_attributes("convergent", "mustprogress", "nounwind", "willreturn", "writeonly")
+            add_fn_attributes("convergent", "mustprogress", "nounwind", "willreturn", "writeonly")
 
         # simd permute
         elseif match(r"air.simd_(ballot|all|vote_all|any|vote_any|shuffle|shuffle_xor|shuffle_down|\
             shuffle_up|shuffle_and_fill_down|shuffle_and_fill_up)", fn) !== nothing
-            add_attributes("convergent", "mustprogress", "nounwind", "willreturn")
+            add_fn_attributes("convergent", "mustprogress", "nounwind", "willreturn")
         end
     end
 

test/metal.jl

@@ -159,6 +159,44 @@ end
     end
 end
 
+@testset "sincos intrinsics" begin
+    mod = @eval module $(gensym())
+        using LLVM.Interop
+
+        function kernel(x)
+            c = Ref{Float32}()
+            s = @typed_ccall("air.sincos.f32", llvmcall, Float32,
+                             (Float32, Ptr{Float32}), x, c)
+            return s + c[]
+        end
+
+        function kernel_fast(x)
+            c = Ref{Float32}()
+            s = @typed_ccall("air.fast_sincos.f32", llvmcall, Float32,
+                             (Float32, Ptr{Float32}), x, c)
+            return s + c[]
+        end
+    end
+
+    # the output argument of sincos should be annotated `nocapture writeonly`
+    # (`nocapture` having been replaced by `captures(none)` in LLVM 21)
+    @test @filecheck begin
+        @check "declare float @air.sincos.f32"
+        @check_same cond=typed_ptrs "(float, float* nocapture writeonly)"
+        @check_same cond=(opaque_ptrs && LLVM.version() < v"21") "(float, ptr nocapture writeonly)"
+        @check_same cond=(LLVM.version() >= v"21") "(float, ptr writeonly captures(none))"
+        Metal.code_llvm(mod.kernel, Tuple{Float32}; dump_module=true)
+    end
+
+    @test @filecheck begin
+        @check "declare float @air.fast_sincos.f32"
+        @check_same cond=typed_ptrs "(float, float* nocapture writeonly)"
+        @check_same cond=(opaque_ptrs && LLVM.version() < v"21") "(float, ptr nocapture writeonly)"
+        @check_same cond=(LLVM.version() >= v"21") "(float, ptr writeonly captures(none))"
+        Metal.code_llvm(mod.kernel_fast, Tuple{Float32}; dump_module=true)
+    end
+end
+
 @testset "unsupported type detection" begin
     mod = @eval module $(gensym())
         function kernel(ptr)