Bug 1750576 - WebGPU shader module asynchronous CompilationInfo API. r=aosmond,emilio

This patch is a lot of plumbing for not that much functionality. The goal is to align CreateShaderModule's error reporting with the spec.
Creating a shader module is now a dedicated async IPDL message returning the compilation info so that it can be exposed as a promise by the WebGPU API.

Differential Revision: https://phabricator.services.mozilla.com/D146817

Cargo.lock

@@ -6180,6 +6180,7 @@ version = "0.1.0"
 dependencies = [
  "bincode",
  "log",
+ "nsstring",
  "parking_lot 0.11.2",
  "serde",
  "wgpu-core",

dom/webgpu/CompilationInfo.cpp

@@ -16,4 +16,19 @@ GPU_IMPL_JS_WRAP(CompilationInfo)
 CompilationInfo::CompilationInfo(ShaderModule* const aParent)
     : ChildOf(aParent) {}
 
+void CompilationInfo::SetMessages(
+    nsTArray<mozilla::webgpu::WebGPUCompilationMessage>& aMessages) {
+  for (auto& msg : aMessages) {
+    mMessages.AppendElement(MakeAndAddRef<mozilla::webgpu::CompilationMessage>(
+        this, msg.lineNum, msg.linePos, msg.offset, std::move(msg.message)));
+  }
+}
+
+void CompilationInfo::GetMessages(
+    nsTArray<RefPtr<mozilla::webgpu::CompilationMessage>>& aMessages) {
+  for (auto& msg : mMessages) {
+    aMessages.AppendElement(msg);
+  }
+}
+
 }  // namespace mozilla::webgpu

dom/webgpu/CompilationInfo.h

@@ -8,6 +8,7 @@
 
 #include "nsWrapperCache.h"
 #include "ObjectModel.h"
+#include "CompilationMessage.h"
 
 namespace mozilla::webgpu {
 class ShaderModule;
@@ -18,10 +19,19 @@ class CompilationInfo final : public nsWrapperCache,
   GPU_DECL_CYCLE_COLLECTION(CompilationInfo)
   GPU_DECL_JS_WRAP(CompilationInfo)
 
- private:
   explicit CompilationInfo(ShaderModule* const aParent);
+
+  void SetMessages(
+      nsTArray<mozilla::webgpu::WebGPUCompilationMessage>& aMessages);
+
+  void GetMessages(
+      nsTArray<RefPtr<mozilla::webgpu::CompilationMessage>>& aMessages);
+
+ private:
   ~CompilationInfo() = default;
   void Cleanup() {}
+
+  nsTArray<RefPtr<mozilla::webgpu::CompilationMessage>> mMessages;
 };
 
 }  // namespace mozilla::webgpu

dom/webgpu/CompilationMessage.cpp

@@ -12,7 +12,13 @@ namespace mozilla::webgpu {
 GPU_IMPL_CYCLE_COLLECTION(CompilationMessage, mParent)
 GPU_IMPL_JS_WRAP(CompilationMessage)
 
-CompilationMessage::CompilationMessage(CompilationInfo* const aParent)
-    : ChildOf(aParent) {}
+CompilationMessage::CompilationMessage(CompilationInfo* const aParent,
+                                       uint64_t aLineNum, uint64_t aLinePos,
+                                       uint64_t aOffset, nsString&& aMessage)
+    : ChildOf(aParent),
+      mLineNum(aLineNum),
+      mLinePos(aLinePos),
+      mOffset(aOffset),
+      mMessage(std::move(aMessage)) {}
 
 }  // namespace mozilla::webgpu

dom/webgpu/CompilationMessage.h

@@ -24,12 +24,19 @@ class CompilationMessage final : public nsWrapperCache,
   uint64_t mLinePos = 0;
   uint64_t mOffset = 0;
   uint64_t mLength = 0;
+  nsString mMessage;
 
  public:
   GPU_DECL_CYCLE_COLLECTION(CompilationMessage)
   GPU_DECL_JS_WRAP(CompilationMessage)
 
-  void GetMessage(dom::DOMString& aMessage) {}
+  explicit CompilationMessage(CompilationInfo* const aParent, uint64_t aLineNum,
+                              uint64_t aLinePos, uint64_t aOffset,
+                              nsString&& aMessage);
+
+  void GetMessage(dom::DOMString& aMessage) {
+    aMessage.AsAString().Assign(mMessage);
+  }
   dom::GPUCompilationMessageType Type() const { return mType; }
   uint64_t LineNum() const { return mLineNum; }
   uint64_t LinePos() const { return mLinePos; }
@@ -37,7 +44,6 @@ class CompilationMessage final : public nsWrapperCache,
   uint64_t Length() const { return mLength; }
 
  private:
-  explicit CompilationMessage(CompilationInfo* const aParent);
   ~CompilationMessage() = default;
   void Cleanup() {}
 };

dom/webgpu/Device.cpp

@@ -266,12 +266,18 @@ already_AddRefed<BindGroup> Device::CreateBindGroup(
 already_AddRefed<ShaderModule> Device::CreateShaderModule(
     JSContext* aCx, const dom::GPUShaderModuleDescriptor& aDesc) {
   Unused << aCx;
-  RawId id = 0;
-  if (mBridge->CanSend()) {
-    id = mBridge->DeviceCreateShaderModule(mId, aDesc);
+
+  if (!mBridge->CanSend()) {
+    return nullptr;
   }
-  RefPtr<ShaderModule> object = new ShaderModule(this, id);
-  return object.forget();
+
+  ErrorResult err;
+  RefPtr<dom::Promise> promise = dom::Promise::Create(GetParentObject(), err);
+  if (NS_WARN_IF(err.Failed())) {
+    return nullptr;
+  }
+
+  return mBridge->DeviceCreateShaderModule(this, aDesc, promise);
 }
 
 already_AddRefed<ComputePipeline> Device::CreateComputePipeline(

dom/webgpu/ShaderModule.cpp

@@ -5,6 +5,7 @@
 
 #include "mozilla/dom/WebGPUBinding.h"
 #include "ShaderModule.h"
+#include "CompilationInfo.h"
 #include "ipc/WebGPUChild.h"
 
 #include "Device.h"
@@ -14,8 +15,9 @@ namespace mozilla::webgpu {
 GPU_IMPL_CYCLE_COLLECTION(ShaderModule, mParent)
 GPU_IMPL_JS_WRAP(ShaderModule)
 
-ShaderModule::ShaderModule(Device* const aParent, RawId aId)
-    : ChildOf(aParent), mId(aId) {}
+ShaderModule::ShaderModule(Device* const aParent, RawId aId,
+                           const RefPtr<dom::Promise>& aCompilationInfo)
+    : ChildOf(aParent), mId(aId), mCompilationInfo(aCompilationInfo) {}
 
 ShaderModule::~ShaderModule() { Cleanup(); }
 
@@ -29,4 +31,9 @@ void ShaderModule::Cleanup() {
   }
 }
 
+already_AddRefed<dom::Promise> ShaderModule::CompilationInfo(ErrorResult& aRv) {
+  RefPtr<dom::Promise> tmp = mCompilationInfo;
+  return tmp.forget();
+}
+
 }  // namespace mozilla::webgpu

dom/webgpu/ShaderModule.h

@@ -12,20 +12,27 @@
 
 namespace mozilla::webgpu {
 
+class CompilationInfo;
 class Device;
 
 class ShaderModule final : public ObjectBase, public ChildOf<Device> {
  public:
-  GPU_DECL_CYCLE_COLLECTION(ShaderModule)
+  GPU_DECL_CYCLE_COLLECTION(
+      ShaderModule)  // TODO: kvark's WIP patch was passing CompilationInfo as a
+                     // second argument here.
   GPU_DECL_JS_WRAP(ShaderModule)
 
-  ShaderModule(Device* const aParent, RawId aId);
+  ShaderModule(Device* const aParent, RawId aId,
+               const RefPtr<dom::Promise>& aCompilationInfo);
+  already_AddRefed<dom::Promise> CompilationInfo(ErrorResult& aRv);
 
   const RawId mId;
 
  private:
   virtual ~ShaderModule();
   void Cleanup();
+
+  RefPtr<dom::Promise> mCompilationInfo;
 };
 
 }  // namespace mozilla::webgpu

dom/webgpu/ipc/PWebGPU.ipdl

@@ -13,6 +13,7 @@ using dom::GPUCommandBufferDescriptor from "mozilla/dom/WebGPUBinding.h";
 using dom::GPUBufferDescriptor from "mozilla/dom/WebGPUBinding.h";
 using webgpu::ffi::WGPUHostMap from "mozilla/webgpu/ffi/wgpu.h";
 using MaybeScopedError from "mozilla/webgpu/WebGPUTypes.h";
+using WebGPUCompilationMessage from "mozilla/webgpu/WebGPUTypes.h";
 
 include "mozilla/ipc/ByteBufUtils.h";
 include "mozilla/layers/LayersMessageUtils.h";
@@ -44,6 +45,8 @@ parent:
   async InstanceRequestAdapter(GPURequestAdapterOptions options, RawId[] ids) returns (ByteBuf byteBuf);
   async AdapterRequestDevice(RawId selfId, ByteBuf buf, RawId newId) returns (bool success);
   async AdapterDestroy(RawId selfId);
+  // TODO: We want to return an array of compilation messages.
+  async DeviceCreateShaderModule(RawId selfId, RawId bufferId, nsString label, nsCString code) returns (WebGPUCompilationMessage[] messages);
   async BufferReturnShmem(RawId selfId, Shmem shmem);
   async BufferMap(RawId selfId, WGPUHostMap hostMap, uint64_t offset, uint64_t size) returns (Shmem sm);
   async BufferUnmap(RawId selfId, Shmem shmem, bool flush, bool keepShmem);

dom/webgpu/ipc/WebGPUChild.cpp

@@ -15,6 +15,7 @@
 #include "Adapter.h"
 #include "DeviceLostInfo.h"
 #include "Sampler.h"
+#include "CompilationInfo.h"
 
 namespace mozilla::webgpu {
 
@@ -705,20 +706,34 @@ RawId WebGPUChild::DeviceCreateBindGroup(
   return id;
 }
 
-RawId WebGPUChild::DeviceCreateShaderModule(
-    RawId aSelfId, const dom::GPUShaderModuleDescriptor& aDesc) {
-  ffi::WGPUShaderModuleDescriptor desc = {};
+already_AddRefed<ShaderModule> WebGPUChild::DeviceCreateShaderModule(
+    Device* aDevice, const dom::GPUShaderModuleDescriptor& aDesc,
+    RefPtr<dom::Promise> aPromise) {
+  RawId deviceId = aDevice->mId;
+  RawId moduleId =
+      ffi::wgpu_client_make_shader_module_id(mClient.get(), deviceId);
 
-  desc.code = reinterpret_cast<const uint8_t*>(aDesc.mCode.get());
-  desc.code_length = aDesc.mCode.Length();
+  RefPtr<ShaderModule> shaderModule =
+      new ShaderModule(aDevice, moduleId, aPromise);
 
-  ByteBuf bb;
-  RawId id = ffi::wgpu_client_create_shader_module(mClient.get(), aSelfId,
-                                                   &desc, ToFFI(&bb));
-  if (!SendDeviceAction(aSelfId, std::move(bb))) {
-    MOZ_CRASH("IPC failure");
-  }
-  return id;
+  nsString noLabel;
+  const nsString& label =
+      aDesc.mLabel.WasPassed() ? aDesc.mLabel.Value() : noLabel;
+  SendDeviceCreateShaderModule(deviceId, moduleId, label, aDesc.mCode)
+      ->Then(
+          GetCurrentSerialEventTarget(), __func__,
+          [aPromise,
+           shaderModule](nsTArray<WebGPUCompilationMessage>&& messages) {
+            RefPtr<CompilationInfo> infoObject(
+                new CompilationInfo(shaderModule));
+            infoObject->SetMessages(messages);
+            aPromise->MaybeResolve(infoObject);
+          },
+          [aPromise](const ipc::ResponseRejectReason& aReason) {
+            aPromise->MaybeRejectWithNotSupportedError("IPC error");
+          });
+
+  return shaderModule.forget();
 }
 
 RawId WebGPUChild::DeviceCreateComputePipelineImpl(

dom/webgpu/ipc/WebGPUChild.h

@@ -83,9 +83,6 @@ class WebGPUChild final : public PWebGPUChild, public SupportsWeakPtr {
       RawId aSelfId, const dom::GPUPipelineLayoutDescriptor& aDesc);
   RawId DeviceCreateBindGroup(RawId aSelfId,
                               const dom::GPUBindGroupDescriptor& aDesc);
-  RawId DeviceCreateShaderModule(RawId aSelfId,
-                                 const dom::GPUShaderModuleDescriptor& aDesc);
-
   RawId DeviceCreateComputePipeline(
       PipelineCreationContext* const aContext,
       const dom::GPUComputePipelineDescriptor& aDesc);
@@ -98,6 +95,9 @@ class WebGPUChild final : public PWebGPUChild, public SupportsWeakPtr {
   RefPtr<PipelinePromise> DeviceCreateRenderPipelineAsync(
       PipelineCreationContext* const aContext,
       const dom::GPURenderPipelineDescriptor& aDesc);
+  already_AddRefed<ShaderModule> DeviceCreateShaderModule(
+      Device* aDevice, const dom::GPUShaderModuleDescriptor& aDesc,
+      RefPtr<dom::Promise> aPromise);
 
   void DeviceCreateSwapChain(RawId aSelfId, const RGBDescriptor& aRgbDesc,
                              size_t maxBufferCount,

dom/webgpu/ipc/WebGPUParent.cpp

@@ -633,6 +633,37 @@ ipc::IPCResult WebGPUParent::RecvDeviceCreateSwapChain(
   return IPC_OK();
 }
 
+ipc::IPCResult WebGPUParent::RecvDeviceCreateShaderModule(
+    RawId aSelfId, RawId aBufferId, const nsString& aLabel,
+    const nsCString& aCode, DeviceCreateShaderModuleResolver&& aOutMessage) {
+  NS_ConvertUTF16toUTF8 label(aLabel);
+
+  ffi::WGPUShaderModuleCompilationMessage message;
+
+  bool ok = ffi::wgpu_server_device_create_shader_module(
+      mContext.get(), aSelfId, aBufferId, label.get(),
+      reinterpret_cast<const uint8_t*>(aCode.get()), aCode.Length(), &message);
+
+  nsTArray<WebGPUCompilationMessage> messages;
+
+  if (!ok) {
+    WebGPUCompilationMessage msg;
+    msg.lineNum = message.line_number;
+    msg.linePos = message.line_pos;
+    msg.offset = message.utf16_offset;
+    msg.length = message.utf16_length;
+    msg.message = message.message;
+    // wgpu currently only returns errors.
+    msg.messageType = WebGPUCompilationMessageType::Error;
+
+    messages.AppendElement(msg);
+  }
+
+  aOutMessage(messages);
+
+  return IPC_OK();
+}
+
 struct PresentRequest {
   const ffi::WGPUGlobal* mContext;
   RefPtr<PresentationData> mData;

dom/webgpu/ipc/WebGPUParent.h

@@ -69,6 +69,10 @@ class WebGPUParent final : public PWebGPUParent {
                                            const layers::RGBDescriptor& aDesc,
                                            const nsTArray<RawId>& aBufferIds,
                                            const CompositableHandle& aHandle);
+  ipc::IPCResult RecvDeviceCreateShaderModule(
+      RawId aSelfId, RawId aModuleId, const nsString& aLabel,
+      const nsCString& aCode, DeviceCreateShaderModuleResolver&& aOutMessage);
+
   ipc::IPCResult RecvSwapChainPresent(const CompositableHandle& aHandle,
                                       RawId aTextureId,
                                       RawId aCommandEncoderId);

dom/webgpu/ipc/WebGPUSerialize.h

@@ -39,6 +39,9 @@ DEFINE_IPC_SERIALIZER_WITH_FIELDS(mozilla::dom::GPUBufferDescriptor, mSize,
 DEFINE_IPC_SERIALIZER_WITH_FIELDS(mozilla::webgpu::ScopedError, operationError,
                                   validationMessage);
 
+DEFINE_IPC_SERIALIZER_WITH_FIELDS(mozilla::webgpu::WebGPUCompilationMessage,
+                                  message, lineNum, linePos);
+
 #undef DEFINE_IPC_SERIALIZER_FFI_ENUM
 #undef DEFINE_IPC_SERIALIZER_DOM_ENUM
 #undef DEFINE_IPC_SERIALIZER_ENUM_GUARD

dom/webgpu/ipc/WebGPUTypes.h

@@ -26,6 +26,22 @@ struct ScopedError {
 };
 using MaybeScopedError = Maybe<ScopedError>;
 
+enum class WebGPUCompilationMessageType { Error, Warning, Info };
+
+// TODO: Better name? CompilationMessage alread taken by the dom object.
+/// The serializable counterpart of the dom object CompilationMessage.
+struct WebGPUCompilationMessage {
+  nsString message;
+  uint64_t lineNum = 0;
+  uint64_t linePos = 0;
+  // In utf16 code units.
+  uint64_t offset = 0;
+  // In utf16 code units.
+  uint64_t length = 0;
+  WebGPUCompilationMessageType messageType =
+      WebGPUCompilationMessageType::Error;
+};
+
 }  // namespace mozilla::webgpu
 
 #endif  // WEBGPU_TYPES_H_

dom/webidl/WebGPU.webidl

@@ -627,9 +627,8 @@ interface GPUCompilationMessage {
 [Pref="dom.webgpu.enabled",
  Exposed=(Window,DedicatedWorker)]
 interface GPUCompilationInfo {
-    //TODO:
-    //[Cached, Frozen, Pure]
-    //readonly attribute sequence<GPUCompilationMessage> messages;
+    [Cached, Frozen, Pure]
+    readonly attribute sequence<GPUCompilationMessage> messages;
 };
 
 // ShaderModule
@@ -643,8 +642,8 @@ dictionary GPUShaderModuleDescriptor : GPUObjectDescriptorBase {
 [Pref="dom.webgpu.enabled",
  Exposed=(Window,DedicatedWorker)]
 interface GPUShaderModule {
-    //TODO:
-    //Promise<GPUCompilationInfo> compilationInfo();
+    [Throws]
+    Promise<GPUCompilationInfo> compilationInfo();
 };
 GPUShaderModule includes GPUObjectBase;
 

gfx/wgpu_bindings/Cargo.toml

@@ -37,3 +37,4 @@ bincode = "1"
 log = "0.4"
 parking_lot = "0.11"
 serde = "1"
+nsstring = { path = "../../xpcom/rust/nsstring" }

gfx/wgpu_bindings/cbindgen.toml

@@ -15,6 +15,8 @@ autogen_warning = """/* DO NOT MODIFY THIS MANUALLY! This file was generated usi
  * - $CBINDGEN is the path to the cbindgen executable provided by mozbuild (the exact version often matters)
  */
 
+#include "nsString.h"
+
 struct WGPUByteBuf;
 typedef uint64_t WGPUNonZeroU64;
 typedef uint64_t WGPUOption_BufferSize;
@@ -27,6 +29,7 @@ typedef uint64_t WGPUOption_BindGroupLayoutId;
 typedef uint64_t WGPUOption_SamplerId;
 typedef uint64_t WGPUOption_SurfaceId;
 typedef uint64_t WGPUOption_TextureViewId;
+typedef nsString WGPUnsString;
 """
 include_version = true
 braces = "SameLine"

gfx/wgpu_bindings/src/client.rs

@@ -35,13 +35,6 @@ fn make_byte_buf<T: serde::Serialize>(data: &T) -> ByteBuf {
     ByteBuf::from_vec(vec)
 }
 
-#[repr(C)]
-pub struct ShaderModuleDescriptor {
-    label: RawString,
-    code: *const u8,
-    code_length: usize,
-}
-
 #[repr(C)]
 pub struct ProgrammableStageDescriptor {
     module: id::ShaderModuleId,
@@ -894,30 +887,17 @@ pub unsafe extern "C" fn wgpu_client_create_bind_group(
 }
 
 #[no_mangle]
-pub unsafe extern "C" fn wgpu_client_create_shader_module(
+pub extern "C" fn wgpu_client_make_shader_module_id(
     client: &Client,
     device_id: id::DeviceId,
-    desc: &ShaderModuleDescriptor,
-    bb: &mut ByteBuf,
 ) -> id::ShaderModuleId {
     let backend = device_id.backend();
-    let id = client
+    client
         .identities
         .lock()
         .select(backend)
         .shader_modules
-        .alloc(backend);
-
-    let code =
-        std::str::from_utf8_unchecked(std::slice::from_raw_parts(desc.code, desc.code_length));
-    let desc = wgc::pipeline::ShaderModuleDescriptor {
-        label: cow_label(&desc.label),
-        shader_bound_checks: wgt::ShaderBoundChecks::new(),
-    };
-
-    let action = DeviceAction::CreateShaderModule(id, desc, Cow::Borrowed(code));
-    *bb = make_byte_buf(&action);
-    id
+        .alloc(backend)
 }
 
 #[no_mangle]

gfx/wgpu_bindings/src/server.rs

@@ -7,10 +7,14 @@ use crate::{
     CommandEncoderAction, DeviceAction, DropAction, QueueWriteAction, RawString, TextureAction,
 };
 
+use nsstring::nsString;
+
 use wgc::{gfx_select, id};
+use wgc::pipeline::CreateShaderModuleError;
 
 use std::sync::atomic::{AtomicU32, Ordering};
 use std::{error::Error, os::raw::c_char, ptr, slice};
+use std::borrow::Cow;
 
 /// A fixed-capacity, null-terminated error buffer owned by C++.
 ///
@@ -207,6 +211,86 @@ pub extern "C" fn wgpu_server_device_drop(global: &Global, self_id: id::DeviceId
     gfx_select!(self_id => global.device_drop(self_id))
 }
 
+impl ShaderModuleCompilationMessage {
+    fn set_error(&mut self, error: &CreateShaderModuleError, source: &str) {
+        // The WebGPU spec says that if the message doesn't point to a particular position in
+        // the source, the line number, position, offset and lengths should be zero.
+        self.line_number = 0;
+        self.line_pos = 0;
+        self.utf16_offset = 0;
+        self.utf16_length = 0;
+
+        if let Some(location) = error.location(source) {
+            self.line_number = location.line_number as u64;
+            self.line_pos = location.line_position as u64;
+
+            let start = location.offset as usize;
+            let end = start + location.length as usize;
+            self.utf16_offset = source[0..start].chars().map(|c| c.len_utf16() as u64).sum();
+            self.utf16_length = source[start..end].chars().map(|c| c.len_utf16() as u64).sum();
+        }
+
+        let error_string = error.to_string();
+
+        if !error_string.is_empty() {
+            self.message = nsString::from(&error_string[..]);
+        }
+    }
+}
+
+/// A compilation message representation for the ffi boundary.
+/// the message is immediately copied into an equivalent C++
+/// structure that owns its strings.
+#[repr(C)]
+#[derive(Clone)]
+pub struct ShaderModuleCompilationMessage {
+    pub line_number: u64,
+    pub line_pos: u64,
+    pub utf16_offset: u64,
+    pub utf16_length: u64,
+    pub message: nsString,
+}
+
+/// Creates a shader module and returns an object describing the errors if any.
+///
+/// If there was no error, the returned pointer is nil.
+#[no_mangle]
+pub extern "C" fn wgpu_server_device_create_shader_module(
+    global: &Global,
+    self_id: id::DeviceId,
+    module_id: id::ShaderModuleId,
+    label: RawString,
+    code: *const u8,
+    code_length: usize,
+    out_message: &mut ShaderModuleCompilationMessage
+) -> bool {
+    unsafe {
+        let label = cow_label(&label);
+
+        let source_str = std::str::from_utf8(std::slice::from_raw_parts(code, code_length)).unwrap();
+        let source = wgc::pipeline::ShaderModuleSource::Wgsl(Cow::from(source_str));
+
+        let desc = wgc::pipeline::ShaderModuleDescriptor {
+            label,
+            shader_bound_checks: wgt::ShaderBoundChecks::new(),
+        };
+
+        let (_, error) = gfx_select!(
+            self_id => global.device_create_shader_module(
+                self_id, &desc, source, module_id
+            )
+        );
+
+        if let Some(err) = error {
+            out_message.set_error(&err, source_str);
+            return false;
+        }
+
+        // Avoid allocating the structure that holds errors in the common case (no errors).
+        return true;
+    }
+}
+
 #[no_mangle]
 pub extern "C" fn wgpu_server_device_create_buffer(
     global: &Global,