[WebGPU] MTLResource allocations should be attributed to the respecti…

…ve web process

https://bugs.webkit.org/show_bug.cgi?id=254219
<radar://107001772>

Reviewed by Tadeu Zagallo and Myles C. Maxfield.

Attribute memory used by IOSurfaces and MTLResource instances to the
corresponding web process, so we do not exceed jetsam limits on iOS.

Mostly this involves piping the web process ID which already is
accessible via RemoteGPU down to WebGPU.framework and then calling the
appropriate IOSurface and Metal SPIs.

* Source/WebCore/PAL/pal/graphics/WebGPU/Impl/WebGPUCreateImpl.cpp:
(PAL::WebGPU::create):
* Source/WebCore/PAL/pal/graphics/WebGPU/Impl/WebGPUCreateImpl.h:
* Source/WebGPU/WebGPU/Buffer.mm:
(WebGPU::Device::safeCreateBuffer const):
* Source/WebGPU/WebGPU/Device.h:
* Source/WebGPU/WebGPU/Device.mm:
(WebGPU::setOwnerWithIdentity):
(WebGPU::Device::setOwnerWithIdentity const):
(WebGPU::Device::newBufferWithBytes const):
(WebGPU::Device::newTextureWithDescriptor const):
* Source/WebGPU/WebGPU/Instance.h:
* Source/WebGPU/WebGPU/Instance.mm:
(WebGPU::Instance::create):
(WebGPU::Instance::Instance):
(WebGPU::m_webProcessID):
(WebGPU::Instance::createSurface):
(WebGPU::Instance::webProcessID const):
* Source/WebGPU/WebGPU/PresentationContext.h:
* Source/WebGPU/WebGPU/PresentationContext.mm:
(WebGPU::PresentationContext::create):
* Source/WebGPU/WebGPU/PresentationContextIOSurface.h:
* Source/WebGPU/WebGPU/PresentationContextIOSurface.mm:
(WebGPU::PresentationContextIOSurface::create):
(WebGPU::PresentationContextIOSurface::PresentationContextIOSurface):
(WebGPU::PresentationContextIOSurface::renderBuffersWereRecreated):
(WebGPU::PresentationContextIOSurface::configure):
* Source/WebGPU/WebGPU/Queue.mm:
(WebGPU::Queue::writeBuffer):
(WebGPU::Queue::writeTexture):
* Source/WebGPU/WebGPU/Texture.mm:
(WebGPU::Device::createTexture):
* Source/WebGPU/WebGPU/WebGPUExt.h:
* Source/WebKit/GPUProcess/graphics/WebGPU/RemoteGPU.cpp:
(WebKit::webProcessID):
(WebKit::RemoteGPU::workQueueInitialize):

* Source/WebGPU/WebGPU/MetalSPI.h:
Add SPI header.

Canonical link: https://commits.webkit.org/276332@main

Source/WebCore/Modules/WebGPU/Implementation/WebGPUCreateImpl.cpp

@@ -43,8 +43,11 @@ WTF_WEAK_LINK_FORCE_IMPORT(wgpuCreateInstance);
 
 namespace WebCore::WebGPU {
 
-RefPtr<GPU> create(ScheduleWorkFunction&& scheduleWorkFunction)
+RefPtr<GPU> create(ScheduleWorkFunction&& scheduleWorkFunction, const WebCore::ProcessIdentity* webProcessIdentity)
 {
+#if !HAVE(TASK_IDENTITY_TOKEN)
+    UNUSED_PARAM(webProcessIdentity);
+#endif
     auto scheduleWorkBlock = makeBlockPtr([scheduleWorkFunction = WTFMove(scheduleWorkFunction)](WGPUWorkItem workItem)
     {
         scheduleWorkFunction(CompletionHandler<void(void)>(makeBlockPtr(WTFMove(workItem)), CompletionHandlerCallThread::AnyThread));
@@ -55,6 +58,11 @@ RefPtr<GPU> create(ScheduleWorkFunction&& scheduleWorkFunction)
             static_cast<WGPUSType>(WGPUSTypeExtended_InstanceCocoaDescriptor),
         },
         scheduleWorkBlock.get(),
+#if HAVE(TASK_IDENTITY_TOKEN)
+        webProcessIdentity ? &webProcessIdentity->taskId() : nullptr,
+#else
+        nullptr,
+#endif
     };
     WGPUInstanceDescriptor descriptor = { &cocoaDescriptor.chain };
 

Source/WebCore/Modules/WebGPU/Implementation/WebGPUCreateImpl.h

@@ -32,11 +32,15 @@
 #include <wtf/Function.h>
 #include <wtf/RefPtr.h>
 
+namespace WebCore {
+class ProcessIdentity;
+}
+
 namespace WebCore::WebGPU {
 
 using WorkItem = CompletionHandler<void(void)>;
 using ScheduleWorkFunction = Function<void(WorkItem&&)>;
-WEBCORE_EXPORT RefPtr<GPU> create(ScheduleWorkFunction&&);
+WEBCORE_EXPORT RefPtr<GPU> create(ScheduleWorkFunction&&, const WebCore::ProcessIdentity*);
 
 } // namespace WebCore::WebGPU
 

Source/WebCore/platform/ProcessIdentity.h

@@ -57,6 +57,7 @@ class ProcessIdentity {
 
 #if HAVE(TASK_IDENTITY_TOKEN)
     task_id_token_t taskIdToken() const { return m_taskIdToken.sendRight(); }
+    const MachSendRight& taskId() const { return m_taskIdToken; }
 #endif
 
 private:

Source/WebGPU/WebGPU.xcodeproj/project.pbxproj

@@ -10,6 +10,7 @@
 		0D30F93729F1F94A0055D9F1 /* ExternalTexture.mm in Sources */ = {isa = PBXBuildFile; fileRef = 0D30F93629F1F94A0055D9F1 /* ExternalTexture.mm */; };
 		0D30F93929F1FAC50055D9F1 /* ExternalTexture.h in Headers */ = {isa = PBXBuildFile; fileRef = 0D30F93829F1FAC50055D9F1 /* ExternalTexture.h */; };
 		0D30F93B29F1FBE40055D9F1 /* CoreVideo.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 0D30F93A29F1FBE40055D9F1 /* CoreVideo.framework */; };
+		0D509DCD29CAB6EC00546D84 /* MetalSPI.h in Headers */ = {isa = PBXBuildFile; fileRef = 0D509DCC29CAB6EC00546D84 /* MetalSPI.h */; };
 		1C0F41EE280940650005886D /* HardwareCapabilities.mm in Sources */ = {isa = PBXBuildFile; fileRef = 1C0F41EC280940650005886D /* HardwareCapabilities.mm */; };
 		1C2CEDEE271E8A7300EDC16F /* Metal.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 1C2CEDED271E8A7300EDC16F /* Metal.framework */; };
 		1C582FF927E04131009B40F0 /* CommandsMixin.mm in Sources */ = {isa = PBXBuildFile; fileRef = 1C582FF727E04131009B40F0 /* CommandsMixin.mm */; };
@@ -260,6 +261,7 @@
 		0D30F93829F1FAC50055D9F1 /* ExternalTexture.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ExternalTexture.h; sourceTree = "<group>"; };
 		0D30F93A29F1FBE40055D9F1 /* CoreVideo.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreVideo.framework; path = System/Library/Frameworks/CoreVideo.framework; sourceTree = SDKROOT; };
 		0D4D2E80294A89CF0000A1AB /* BindableResource.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = BindableResource.h; sourceTree = "<group>"; };
+		0D509DCC29CAB6EC00546D84 /* MetalSPI.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = MetalSPI.h; sourceTree = "<group>"; };
 		1C0F41EC280940650005886D /* HardwareCapabilities.mm */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.objcpp; path = HardwareCapabilities.mm; sourceTree = "<group>"; };
 		1C0F41ED280940650005886D /* HardwareCapabilities.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = HardwareCapabilities.h; sourceTree = "<group>"; };
 		1C2CEDED271E8A7300EDC16F /* Metal.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = Metal.framework; path = System/Library/Frameworks/Metal.framework; sourceTree = SDKROOT; };
@@ -581,6 +583,7 @@
 				1C5ACAA0273A426D0095F8D5 /* Instance.h */,
 				1C5ACA92273A41C20095F8D5 /* Instance.mm */,
 				1C33755D27FA23B8002F1644 /* IsValidToUseWith.h */,
+				0D509DCC29CAB6EC00546D84 /* MetalSPI.h */,
 				973F784529C8A78200166C66 /* Pipeline.h */,
 				973F784629C8A78200166C66 /* Pipeline.mm */,
 				1C5ACADA273A4E710095F8D5 /* PipelineLayout.h */,
@@ -828,6 +831,7 @@
 			files = (
 				1CEBD7E72716AFBA00A5254D /* WebGPU.h in Headers */,
 				0D30F93929F1FAC50055D9F1 /* ExternalTexture.h in Headers */,
+				0D509DCD29CAB6EC00546D84 /* MetalSPI.h in Headers */,
 				973F784729C8A78200166C66 /* Pipeline.h in Headers */,
 				1C5ACAD3273A4C860095F8D5 /* WebGPUExt.h in Headers */,
 			);

Source/WebGPU/WebGPU/Buffer.mm

@@ -99,9 +99,9 @@ static MTLStorageMode storageMode(bool deviceHasUnifiedMemory, WGPUBufferUsageFl
 id<MTLBuffer> Device::safeCreateBuffer(NSUInteger length, MTLStorageMode storageMode, MTLCPUCacheMode cpuCacheMode, MTLHazardTrackingMode hazardTrackingMode) const
 {
     MTLResourceOptions resourceOptions = (cpuCacheMode << MTLResourceCPUCacheModeShift) | (storageMode << MTLResourceStorageModeShift) | (hazardTrackingMode << MTLResourceHazardTrackingModeShift);
-    // FIXME(PERFORMANCE): Consider returning nil instead of clamping to 1.
-    // FIXME(PERFORMANCE): Suballocate multiple Buffers either from MTLHeaps or from larger MTLBuffers.
-    return [m_device newBufferWithLength:std::max(static_cast<NSUInteger>(1), length) options:resourceOptions];
+    id<MTLBuffer> buffer = [m_device newBufferWithLength:std::max<NSUInteger>(1, length) options:resourceOptions];
+    setOwnerWithIdentity(buffer);
+    return buffer;
 }
 
 Ref<Buffer> Device::createBuffer(const WGPUBufferDescriptor& descriptor)

Source/WebGPU/WebGPU/Device.h

@@ -125,6 +125,9 @@ class Device : public WGPUDeviceImpl, public ThreadSafeRefCountedAndCanMakeThrea
     uint32_t maxBuffersForFragmentStage() const;
     uint32_t maxBuffersForComputeStage() const;
     uint32_t vertexBufferIndexForBindGroup(uint32_t groupIndex) const;
+    id<MTLBuffer> newBufferWithBytes(const void*, size_t, MTLResourceOptions) const;
+    id<MTLBuffer> newBufferWithBytesNoCopy(void*, size_t, MTLResourceOptions) const;
+    id<MTLTexture> newTextureWithDescriptor(MTLTextureDescriptor *, IOSurfaceRef = nullptr, NSUInteger plane = 0) const;
 
     static bool isStencilOnlyFormat(MTLPixelFormat);
     bool shouldStopCaptureAfterSubmit();
@@ -158,6 +161,7 @@ class Device : public WGPUDeviceImpl, public ThreadSafeRefCountedAndCanMakeThrea
         simd::float4x3 colorSpaceConversionMatrix;
     };
     ExternalTextureData createExternalTextureFromPixelBuffer(CVPixelBufferRef, WGPUColorSpace) const;
+    void setOwnerWithIdentity(id<MTLResource>) const;
 
     struct Error {
         WGPUErrorType type;

Source/WebGPU/WebGPU/Device.mm

@@ -33,6 +33,7 @@
 #import "Buffer.h"
 #import "CommandEncoder.h"
 #import "ComputePipeline.h"
+#import "MetalSPI.h"
 #import "PipelineLayout.h"
 #import "PresentationContext.h"
 #import "QuerySet.h"
@@ -237,6 +238,28 @@ static void captureFrame(id<MTLDevice> captureObject)
     m_isLost = true;
 }
 
+static void setOwnerWithIdentity(id<MTLResourceSPI> resource, auto webProcessID)
+{
+    if (!resource)
+        return;
+
+    if (![resource respondsToSelector:@selector(setOwnerWithIdentity:)])
+        return;
+
+    [resource setOwnerWithIdentity:webProcessID];
+}
+
+void Device::setOwnerWithIdentity(id<MTLResource> resource) const
+{
+    if (auto optionalWebProcessID = instance().webProcessID()) {
+        auto webProcessID = optionalWebProcessID->sendRight();
+        if (!webProcessID)
+            return;
+
+        WebGPU::setOwnerWithIdentity((id<MTLResourceSPI>)resource, webProcessID);
+    }
+}
+
 void Device::destroy()
 {
     m_destroyed = true;
@@ -369,6 +392,27 @@ static void captureFrame(id<MTLDevice> captureObject)
     return WGSL::vertexBufferIndexForBindGroup(groupIndex, maxBuffersPlusVertexBuffersForVertexStage() - 1);
 }
 
+id<MTLBuffer> Device::newBufferWithBytes(const void* pointer, size_t length, MTLResourceOptions options) const
+{
+    id<MTLBuffer> buffer = [m_device newBufferWithBytes:pointer length:length options:options];
+    setOwnerWithIdentity(buffer);
+    return buffer;
+}
+
+id<MTLBuffer> Device::newBufferWithBytesNoCopy(void* pointer, size_t length, MTLResourceOptions options) const
+{
+    id<MTLBuffer> buffer = [m_device newBufferWithBytesNoCopy:pointer length:length options:options deallocator:nil];
+    setOwnerWithIdentity(buffer);
+    return buffer;
+}
+
+id<MTLTexture> Device::newTextureWithDescriptor(MTLTextureDescriptor *textureDescriptor, IOSurfaceRef ioSurface, NSUInteger plane) const
+{
+    id<MTLTexture> texture = ioSurface ? [m_device newTextureWithDescriptor:textureDescriptor iosurface:ioSurface plane:plane] : [m_device newTextureWithDescriptor:textureDescriptor];
+    setOwnerWithIdentity(texture);
+    return texture;
+}
+
 void Device::captureFrameIfNeeded() const
 {
     GPUFrameCapture::captureSingleFrameIfNeeded(m_device);

Source/WebGPU/WebGPU/Instance.h

@@ -29,12 +29,17 @@
 #import <wtf/Deque.h>
 #import <wtf/FastMalloc.h>
 #import <wtf/Lock.h>
+#import <wtf/MachSendRight.h>
 #import <wtf/Ref.h>
 #import <wtf/ThreadSafeRefCounted.h>
 
 struct WGPUInstanceImpl {
 };
 
+namespace WTF {
+class MachSendRight;
+}
+
 namespace WebGPU {
 
 class Adapter;
@@ -61,16 +66,18 @@ class Instance : public WGPUInstanceImpl, public ThreadSafeRefCounted<Instance>
     // This can be called on a background thread.
     using WorkItem = CompletionHandler<void(void)>;
     void scheduleWork(WorkItem&&);
+    const std::optional<const MachSendRight>& webProcessID() const;
 
 private:
-    Instance(WGPUScheduleWorkBlock);
+    Instance(WGPUScheduleWorkBlock, const WTF::MachSendRight* webProcessResourceOwner);
     explicit Instance();
 
     // This can be called on a background thread.
     void defaultScheduleWork(WGPUWorkItem&&);
 
     // This can be used on a background thread.
     Deque<WGPUWorkItem> m_pendingWork WTF_GUARDED_BY_LOCK(m_lock);
+    const std::optional<const MachSendRight> m_webProcessID;
     const WGPUScheduleWorkBlock m_scheduleWorkBlock;
     Lock m_lock;
     bool m_isValid { true };

Source/WebGPU/WebGPU/Instance.mm

@@ -33,14 +33,15 @@
 #import <cstring>
 #import <dlfcn.h>
 #import <wtf/BlockPtr.h>
+#import <wtf/MachSendRight.h>
 #import <wtf/StdLibExtras.h>
 
 namespace WebGPU {
 
 Ref<Instance> Instance::create(const WGPUInstanceDescriptor& descriptor)
 {
     if (!descriptor.nextInChain)
-        return adoptRef(*new Instance(nullptr));
+        return Instance::createInvalid();
 
     if (descriptor.nextInChain->sType != static_cast<WGPUSType>(WGPUSTypeExtended_InstanceCocoaDescriptor))
         return Instance::createInvalid();
@@ -50,11 +51,12 @@
     if (cocoaDescriptor.chain.next)
         return Instance::createInvalid();
 
-    return adoptRef(*new Instance(cocoaDescriptor.scheduleWorkBlock));
+    return adoptRef(*new Instance(cocoaDescriptor.scheduleWorkBlock, reinterpret_cast<const WTF::MachSendRight*>(cocoaDescriptor.webProcessResourceOwner)));
 }
 
-Instance::Instance(WGPUScheduleWorkBlock scheduleWorkBlock)
-    : m_scheduleWorkBlock(scheduleWorkBlock ? WTFMove(scheduleWorkBlock) : ^(WGPUWorkItem workItem) { defaultScheduleWork(WTFMove(workItem)); })
+Instance::Instance(WGPUScheduleWorkBlock scheduleWorkBlock, const MachSendRight* webProcessResourceOwner)
+    : m_webProcessID(webProcessResourceOwner ? std::optional<MachSendRight>(*webProcessResourceOwner) : std::nullopt)
+    , m_scheduleWorkBlock(scheduleWorkBlock ? WTFMove(scheduleWorkBlock) : ^(WGPUWorkItem workItem) { defaultScheduleWork(WTFMove(workItem)); })
 {
 }
 
@@ -68,14 +70,19 @@
 
 Ref<PresentationContext> Instance::createSurface(const WGPUSurfaceDescriptor& descriptor)
 {
-    return PresentationContext::create(descriptor);
+    return PresentationContext::create(descriptor, *this);
 }
 
 void Instance::scheduleWork(WorkItem&& workItem)
 {
     m_scheduleWorkBlock(makeBlockPtr(WTFMove(workItem)).get());
 }
 
+const std::optional<const MachSendRight>& Instance::webProcessID() const
+{
+    return m_webProcessID;
+}
+
 void Instance::defaultScheduleWork(WGPUWorkItem&& workItem)
 {
     Locker locker(m_lock);

Source/WebGPU/WebGPU/MetalSPI.h

@@ -0,0 +1,33 @@
+/*
+ * Copyright (c) 2023 Apple Inc. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ * 1. Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY APPLE INC. ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL APPLE INC. OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#if USE(INTERNAL_APPLE_SDK)
+#import <Metal/MTLResource_Private.h>
+#else
+@protocol MTLResourceSPI <MTLResource>
+@optional
+- (kern_return_t)setOwnerWithIdentity:(mach_port_t)task_id_token;
+@end
+#endif

Source/WebGPU/WebGPU/PresentationContext.h

@@ -41,13 +41,14 @@ namespace WebGPU {
 
 class Adapter;
 class Device;
+class Instance;
 class Texture;
 class TextureView;
 
 class PresentationContext : public WGPUSurfaceImpl, public WGPUSwapChainImpl, public RefCounted<PresentationContext> {
     WTF_MAKE_FAST_ALLOCATED;
 public:
-    static Ref<PresentationContext> create(const WGPUSurfaceDescriptor&);
+    static Ref<PresentationContext> create(const WGPUSurfaceDescriptor&, const Instance&);
     static Ref<PresentationContext> createInvalid()
     {
         return adoptRef(*new PresentationContext());

Source/WebGPU/WebGPU/PresentationContext.mm

@@ -39,14 +39,14 @@
     return presentationContext;
 }
 
-Ref<PresentationContext> PresentationContext::create(const WGPUSurfaceDescriptor& descriptor)
+Ref<PresentationContext> PresentationContext::create(const WGPUSurfaceDescriptor& descriptor, const Instance& instance)
 {
     if (!descriptor.nextInChain || descriptor.nextInChain->next)
         return PresentationContext::createInvalid();
 
     switch (static_cast<unsigned>(descriptor.nextInChain->sType)) {
     case WGPUSTypeExtended_SurfaceDescriptorCocoaSurfaceBacking:
-        return PresentationContextIOSurface::create(descriptor);
+        return PresentationContextIOSurface::create(descriptor, instance);
     case WGPUSType_SurfaceDescriptorFromMetalLayer:
         return PresentationContextCoreAnimation::create(descriptor);
     default:

Source/WebGPU/WebGPU/PresentationContextIOSurface.h

@@ -26,17 +26,19 @@
 #pragma once
 
 #import "PresentationContext.h"
+#import <wtf/MachSendRight.h>
 #import <wtf/Vector.h>
 #import <wtf/spi/cocoa/IOSurfaceSPI.h>
 
 namespace WebGPU {
 
 class Device;
+class Instance;
 
 class PresentationContextIOSurface : public PresentationContext {
     WTF_MAKE_FAST_ALLOCATED;
 public:
-    static Ref<PresentationContextIOSurface> create(const WGPUSurfaceDescriptor&);
+    static Ref<PresentationContextIOSurface> create(const WGPUSurfaceDescriptor&, const Instance&);
 
     virtual ~PresentationContextIOSurface();
 
@@ -50,7 +52,7 @@ class PresentationContextIOSurface : public PresentationContext {
     bool isPresentationContextIOSurface() const override { return true; }
 
 private:
-    PresentationContextIOSurface(const WGPUSurfaceDescriptor&);
+    PresentationContextIOSurface(const WGPUSurfaceDescriptor&, const Instance&);
 
     void renderBuffersWereRecreated(NSArray<IOSurface *> *renderBuffers);
     void onSubmittedWorkScheduled(CompletionHandler<void()>&&);
@@ -66,6 +68,9 @@ class PresentationContextIOSurface : public PresentationContext {
     size_t m_currentIndex { 0 };
     id<MTLFunction> m_luminanceClampFunction;
     id<MTLComputePipelineState> m_computePipelineState;
+#if HAVE(IOSURFACE_SET_OWNERSHIP_IDENTITY) && HAVE(TASK_IDENTITY_TOKEN)
+    std::optional<const MachSendRight> m_webProcessID;
+#endif
 };
 
 } // namespace WebGPU