Skip to content

Commit 7a00e02

Browse files
christiangnrdchristiangnrd
authored
Add MTLResidencySet (JuliaGPU#686)
1 parent 0e59fdb commit 7a00e02

File tree

6 files changed

+165
-3
lines changed

6 files changed

+165
-3
lines changed

‎lib/mtl/MTL.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ include("events.jl")
3434
include("fences.jl")
3535
include("heap.jl")
3636
include("buffer.jl")
37+
include("residency_set.jl")
3738
include("command_queue.jl")
3839
include("command_buf.jl")
3940
include("compute_pipeline.jl")

‎lib/mtl/command_buf.jl‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,15 @@ function commit!(f::Base.Callable, cmdbuf::MTLCommandBuffer)
6868
return ret
6969
end
7070

71+
function use_residency_set!(cmdbuf::MTLCommandBuffer, resset::MTLResidencySet)
72+
@objc [cmdbuf::id{MTLCommandBuffer} useResidencySet:resset::id{MTLResidencySet}]::Nothing
73+
end
74+
75+
function use_residency_sets!(cmdbuf::MTLCommandBuffer, ressets, count)
76+
@objc [cmdbuf::id{MTLCommandBuffer} useResidencySets:resset::Ptr{id{MTLResidencySet}}
77+
count:count::NSUInteger]::Nothing
78+
end
79+
7180
"""
7281
wait_scheduled(commandBuffer)
7382

‎lib/mtl/libmtl.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3294,7 +3294,7 @@ end
32943294
@autoproperty initialCapacity::UInt64 setter = setInitialCapacity
32953295
end
32963296

3297-
@objcwrapper immutable = true availability = macos(v"15.0.0") MTLResidencySet <: NSObject
3297+
@objcwrapper immutable = false availability = macos(v"15.0.0") MTLResidencySet <: NSObject
32983298

32993299
@objcproperties MTLResidencySet begin
33003300
@autoproperty device::id{MTLDevice}

‎lib/mtl/residency_set.jl‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
export MTLResidencySet, MTLResidencySetDescriptor
2+
3+
function MTLResidencySetDescriptor()
4+
desc = @objc [MTLResidencySetDescriptor alloc]::id{MTLResidencySetDescriptor}
5+
obj = MTLResidencySetDescriptor(desc)
6+
return obj
7+
end
8+
9+
function MTLResidencySet(device::MTLDevice, desc::MTLResidencySetDescriptor)
10+
err = Ref{id{NSError}}(nil)
11+
handle = @objc [device::id{MTLDevice} newResidencySetWithDescriptor:desc::id{MTLResidencySetDescriptor}
12+
error:err::Ptr{id{NSError}}]::id{MTLResidencySet}
13+
err[] == nil || throw(NSError(err[]))
14+
obj = MTLResidencySet(handle)
15+
finalizer(release, obj)
16+
return obj
17+
end
18+
19+
# Buffer Arguments
20+
function add_allocation!(resset::MTLResidencySet, allocation)
21+
@objc [resset::id{MTLResidencySet} addAllocation:allocation::id{MTLAllocation}]::Nothing
22+
end
23+
function add_allocations!(resset::MTLResidencySet, allocations, count=length(allocations))
24+
@objc [resset::id{MTLResidencySet} addAllocations:allocations::Ptr{id{MTLAllocation}}
25+
count:count::NSUInteger]::Nothing
26+
end
27+
28+
function remove_all_allocations!(resset::MTLResidencySet)
29+
@objc [resset::id{MTLResidencySet} removeAllAllocations]::Nothing
30+
end
31+
function remove_allocation!(resset::MTLResidencySet, allocation)
32+
@objc [resset::id{MTLResidencySet} removeAllocation:allocation::id{MTLAllocation}]::Nothing
33+
end
34+
function remove_allocations!(resset::MTLResidencySet, allocations, count=length(allocations))
35+
@objc [resset::id{MTLResidencySet} removeAllocations:allocations::Ptr{id{MTLAllocation}}
36+
count:count::NSUInteger]::Nothing
37+
end
38+
39+
function commit!(resset::MTLResidencySet)
40+
@objc [resset::id{MTLResidencySet} commit]::Nothing
41+
end
42+
function request_residency!(resset::MTLResidencySet)
43+
@objc [resset::id{MTLResidencySet} requestResidency]::Nothing
44+
end
45+
function end_residency!(resset::MTLResidencySet)
46+
@objc [resset::id{MTLResidencySet} endResidency]::Nothing
47+
end
48+
function contains_allocation(resset::MTLResidencySet, allocation)
49+
@objc [resset::id{MTLResidencySet} containsAllocation:allocation::id{MTLAllocation}]::Bool
50+
end

‎res/wrap/libmtl.toml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,9 @@ immutable=false
106106
[api.MTLLibrary]
107107
immutable=false
108108

109+
[api.MTLResidencySet]
110+
immutable=false
111+
109112
[api.MTLSharedEvent]
110113
immutable=false
111114

‎test/mtl/metal.jl‎

Lines changed: 101 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -462,13 +462,13 @@ end
462462

463463

464464
Metal.encode_wait!(buf2, event, signal_value)
465-
Metal.commit!(buf2)
465+
commit!(buf2)
466466

467467
unsafe_copyto!(dev, pointer(a), pointer(B), N, queue=queue1, async=true) # GPU -> CPU
468468
unsafe_copyto!(dev, pointer(A), pointer(a), N, queue=queue1, async=true) # CPU -> GPU
469469

470470
Metal.encode_signal!(buf1, event, signal_value)
471-
Metal.commit!(buf1)
471+
commit!(buf1)
472472

473473

474474
Metal.wait_completed(buf2)
@@ -489,6 +489,105 @@ end
489489
end
490490
end
491491

492+
if Metal.is_macos(v"15")
493+
@testset "residency sets" begin
494+
495+
dev = first(devices())
496+
497+
let desc = MTLResidencySetDescriptor()
498+
@test isnothing(desc.label)
499+
desc.label = "Test MTLResidencySetDescriptor"
500+
@test desc.label == "Test MTLResidencySetDescriptor"
501+
502+
@test desc.initialCapacity == 0
503+
desc.initialCapacity = 2
504+
@test desc.initialCapacity == 2
505+
end
506+
507+
desc = MTLResidencySetDescriptor()
508+
desc.initialCapacity = 2
509+
let rset = MTLResidencySet(dev, desc)
510+
alloc1 = MTL.MTLBuffer(dev, 400)
511+
alloc2 = MTL.MTLBuffer(dev, 1024)
512+
allocs = [alloc1, alloc2]
513+
514+
# default parameters
515+
516+
@test isnothing(rset.label)
517+
@test rset.device == dev
518+
@test rset.allocationCount == 0
519+
@test rset.allocatedSize == 0
520+
@test isempty(rset.allAllocations)
521+
@test rset.retainCount == 1
522+
@test contains(rset.description, "ResidencySet")
523+
@test !MTL.contains_allocation(rset, alloc1)
524+
# @test contains(String(rset.description), "ResidencySet")
525+
526+
# one allocation
527+
MTL.add_allocation!(rset, alloc1)
528+
@test MTL.contains_allocation(rset, alloc1)
529+
@test alloc1 in rset.allAllocations
530+
@test rset.allocationCount == 1
531+
@test rset.allocatedSize == 0 # Updated after it's committed
532+
commit!(rset)
533+
@test rset.allocationCount == 1
534+
@test rset.allocatedSize > 0
535+
MTL.remove_allocation!(rset, alloc1)
536+
@test !MTL.contains_allocation(rset, alloc1)
537+
@test rset.allocationCount == 1
538+
@test rset.allocatedSize > 0
539+
commit!(rset)
540+
@test rset.allocationCount == 0
541+
@test rset.allocatedSize == 0
542+
543+
# two allocations
544+
545+
MTL.add_allocations!(rset, allocs)
546+
@test MTL.contains_allocation(rset, alloc2)
547+
@test MTL.contains_allocation(rset, alloc2)
548+
@test alloc1 in rset.allAllocations
549+
@test alloc2 in rset.allAllocations
550+
@test rset.allocationCount == 2
551+
@test rset.allocatedSize == 0 # Updated after it's committed
552+
commit!(rset)
553+
@test rset.allocationCount == 2
554+
@test rset.allocatedSize > 0
555+
MTL.remove_allocations!(rset, allocs)
556+
@test !MTL.contains_allocation(rset, alloc1)
557+
@test !MTL.contains_allocation(rset, alloc2)
558+
@test rset.allocationCount == 2
559+
@test rset.allocatedSize > 0
560+
commit!(rset)
561+
@test rset.allocationCount == 0
562+
@test rset.allocatedSize == 0
563+
564+
# two allocations but different
565+
566+
MTL.add_allocations!(rset, allocs, 2)
567+
@test MTL.contains_allocation(rset, alloc2)
568+
@test MTL.contains_allocation(rset, alloc2)
569+
@test alloc1 in rset.allAllocations
570+
@test alloc2 in rset.allAllocations
571+
@test rset.allocationCount == 2
572+
@test rset.allocatedSize == 0 # Updated after it's committed
573+
commit!(rset)
574+
MTL.request_residency!(rset) # ensure no crash
575+
@test rset.allocationCount == 2
576+
@test rset.allocatedSize > 0
577+
MTL.end_residency!(rset) # ensure no crash
578+
MTL.remove_all_allocations!(rset)
579+
@test !MTL.contains_allocation(rset, alloc1)
580+
@test !MTL.contains_allocation(rset, alloc2)
581+
@test rset.allocationCount == 2
582+
@test rset.allocatedSize > 0
583+
commit!(rset)
584+
@test rset.allocationCount == 0
585+
@test rset.allocatedSize == 0
586+
end
587+
588+
end # @testset "residency sets" begin
589+
end # if Metal.is_macos(v"15")
590+
492591
# TODO: continue adding tests
493592

494593
end # @autoreleasepool begin

0 commit comments

Comments
 (0)