allow creating cuda IPC handles for all Memory types except PinnedMemory and ManagedMemory (#442)

Author: trxcllnt

modules/cuda/src/buffer.ts

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2021, NVIDIA CORPORATION.
+// Copyright (c) 2020-2022, NVIDIA CORPORATION.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
 
 import {Math, runtime} from './addon';
 import {BigIntArray, MemoryData, TypedArray, TypedArrayConstructor} from './interfaces';
-import {DeviceMemory, IpcHandle, Memory} from './memory';
+import {DeviceMemory, IpcHandle, ManagedMemory, Memory, PinnedMemory} from './memory';
 import {
   clampRange,
   isArrayBufferLike,
@@ -319,10 +319,15 @@ export abstract class MemoryView<T extends TypedArray|BigIntArray = any> impleme
    * @summary Create an IpcHandle for the underlying CUDA device memory.
    */
   public getIpcHandle() {
-    if (!(this.buffer instanceof DeviceMemory)) {
-      throw new Error(`${this[Symbol.toStringTag]}'s buffer must be an instance of DeviceMemory`);
+    if (this.buffer instanceof PinnedMemory) {
+      throw new Error(
+        `${this[Symbol.toStringTag]}'s buffer must not be an instance of PinnedMemory`);
     }
-    return new IpcHandle(this.buffer, this.byteOffset);
+    if (this.buffer instanceof ManagedMemory) {
+      throw new Error(
+        `${this[Symbol.toStringTag]}'s buffer must not be an instance of ManagedMemory`);
+    }
+    return new IpcHandle(<any>this.buffer, this.byteOffset);
   }
 }
 

modules/cuda/src/memory.ts

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2021, NVIDIA CORPORATION.
+// Copyright (c) 2020-2022, NVIDIA CORPORATION.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -54,14 +54,16 @@ export class IpcHandle extends CUDAIpcHandle {
 
   /**
    * @summary An object describing the exported {@link DeviceMemory} and CUDA IPC handle.
-   * @returns An object with the device ordinal, byte offset (if applicable) into the
-   *   exported {@link DeviceMemory}, and the 64-bit IPC handle (as a JavaScript Array of octets).
+   * @returns An object with the device ordinal, the 64-bit IPC handle (as a JavaScript Array of
+   *   octets), byte offset (if applicable) into the exported {@link DeviceMemory}, byte length of
+   *   the IPC segment.
    */
   public toJSON() {
     return {
       device: this.device,
-      byteOffset: this.byteOffset,
       handle: [...this.handle],
+      byteOffset: this.byteOffset,
+      byteLength: this.buffer.byteLength - this.byteOffset,
     };
   }
 }

modules/cuda/src/memory/memory.cpp

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2021, NVIDIA CORPORATION.
+// Copyright (c) 2020-2022, NVIDIA CORPORATION.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -69,6 +69,19 @@ Napi::Value cudaMemGetInfoNapi(CallbackArgs const& args) {
                          std::vector<std::string>{"free", "total"});
 }
 
+Napi::Value cudaPointerGetAttributesNapi(CallbackArgs const& args) {
+  auto env   = args.Env();
+  void* dptr = args[0];
+  CUDARTAPI::cudaPointerAttributes attrs{};
+  NODE_CUDA_TRY(CUDARTAPI::cudaPointerGetAttributes(&attrs, dptr), env);
+  auto obj = Napi::Object::New(env);
+  obj.Set("type", attrs.type);
+  obj.Set("device", attrs.device);
+  obj.Set("hptr", reinterpret_cast<uint64_t>(attrs.hostPointer));
+  obj.Set("dptr", reinterpret_cast<uint64_t>(attrs.devicePointer));
+  return obj;
+}
+
 // CUresult cuPointerGetAttribute(void *data, CUpointer_attribute attribute,
 // CUdeviceptr ptr);
 Napi::Value cuPointerGetAttributeNapi(CallbackArgs const& args) {
@@ -140,6 +153,7 @@ Napi::Object initModule(Napi::Env const& env,
   EXPORT_FUNC(env, runtime, "cudaMemset", cudaMemsetNapi);
   EXPORT_FUNC(env, runtime, "cudaMemcpy", cudaMemcpyNapi);
   EXPORT_FUNC(env, runtime, "cudaMemGetInfo", cudaMemGetInfoNapi);
+  EXPORT_FUNC(env, runtime, "cudaPointerGetAttributes", cudaPointerGetAttributesNapi);
   EXPORT_FUNC(env, driver, "cuPointerGetAttribute", cuPointerGetAttributeNapi);
 
   auto PointerAttributes = Napi::Object::New(env);

modules/cuda/src/node_cuda/utilities/cpp_to_napi.hpp

@@ -99,3 +99,12 @@ inline Napi::Value CPPToNapi::operator()(cudaDeviceProp const& props) const {
 }
 
 }  // namespace nv
+
+namespace Napi {
+
+template <>
+inline Value Value::From(napi_env env, CUDARTAPI::cudaMemoryType const& type) {
+  return Value::From(env, static_cast<uint8_t>(type));
+}
+
+}  // namespace Napi