Improve OpenReg test coverage (pytorch#167819)

- add failure-path tests for device, stream, memory, event APIs
- cover async memcpy, pointer attributes, event timing, addTask errors
- verified via cmake --build build && ctest --test-dir build

Pull Request resolved: pytorch#167819
Approved by: https://github.com/fffrog

Author: hipudding

test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/tests/device_tests.cpp

@@ -16,12 +16,22 @@ TEST_F(DeviceTest, GetDeviceCountValid) {
   EXPECT_EQ(count, 2);
 }
 
+TEST_F(DeviceTest, GetDeviceCountNullptr) {
+  // orGetDeviceCount should reject null output pointers.
+  EXPECT_EQ(orGetDeviceCount(nullptr), orErrorUnknown);
+}
+
 TEST_F(DeviceTest, GetDeviceValid) {
   int device = -1;
   EXPECT_EQ(orGetDevice(&device), orSuccess);
   EXPECT_EQ(device, 0);
 }
 
+TEST_F(DeviceTest, GetDeviceNullptr) {
+  // Defensive path: null output pointer must return an error.
+  EXPECT_EQ(orGetDevice(nullptr), orErrorUnknown);
+}
+
 TEST_F(DeviceTest, SetDeviceValid) {
   EXPECT_EQ(orSetDevice(1), orSuccess);
 
@@ -38,4 +48,9 @@ TEST_F(DeviceTest, SetDeviceInvalidNegative) {
   EXPECT_EQ(orSetDevice(-1), orErrorUnknown);
 }
 
+TEST_F(DeviceTest, SetDeviceInvalidTooLarge) {
+  // Device indices are 0-based and strictly less than DEVICE_COUNT (2).
+  EXPECT_EQ(orSetDevice(2), orErrorUnknown);
+}
+
 } // namespace

test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/tests/event_tests.cpp

@@ -29,6 +29,13 @@ TEST_F(EventTest, EventCreateWithFlagsTiming) {
   EXPECT_EQ(orEventDestroy(event), orSuccess);
 }
 
+TEST_F(EventTest, EventCreationNullptr) {
+  // Creation APIs must fail fast on null handles to mirror CUDA semantics.
+  EXPECT_EQ(orEventCreate(nullptr), orErrorUnknown);
+  EXPECT_EQ(
+      orEventCreateWithFlags(nullptr, orEventEnableTiming), orErrorUnknown);
+}
+
 TEST_F(EventTest, EventRecordAndSynchronize) {
   orStream_t stream = nullptr;
   EXPECT_EQ(orStreamCreate(&stream), orSuccess);
@@ -44,6 +51,23 @@ TEST_F(EventTest, EventRecordAndSynchronize) {
   EXPECT_EQ(orStreamDestroy(stream), orSuccess);
 }
 
+TEST_F(EventTest, EventRecordInvalidArgs) {
+  orEvent_t event = nullptr;
+  EXPECT_EQ(orEventCreate(&event), orSuccess);
+
+  orStream_t stream = nullptr;
+  EXPECT_EQ(orStreamCreate(&stream), orSuccess);
+
+  // Record/sync/destroy should validate both stream and event pointers.
+  EXPECT_EQ(orEventRecord(nullptr, stream), orErrorUnknown);
+  EXPECT_EQ(orEventRecord(event, nullptr), orErrorUnknown);
+  EXPECT_EQ(orEventSynchronize(nullptr), orErrorUnknown);
+  EXPECT_EQ(orEventDestroy(nullptr), orErrorUnknown);
+
+  EXPECT_EQ(orEventDestroy(event), orSuccess);
+  EXPECT_EQ(orStreamDestroy(stream), orSuccess);
+}
+
 TEST_F(EventTest, EventElapsedTime) {
   orStream_t stream = nullptr;
   EXPECT_EQ(orStreamCreate(&stream), orSuccess);
@@ -70,6 +94,60 @@ TEST_F(EventTest, EventElapsedTime) {
   EXPECT_EQ(orEventDestroy(end), orSuccess);
 }
 
+// TODO: recording events to a stream is not allowed
+// if the stream and the event are not on the same device
+// Uncomment this test case after the issue is fixed.
+// see #167819
+TEST_F(EventTest, DISABLED_EventElapsedTimeDifferentDevicesFails) {
+  orStream_t stream = nullptr;
+  EXPECT_EQ(orStreamCreate(&stream), orSuccess);
+
+  orEvent_t start = nullptr;
+  orEvent_t end = nullptr;
+  EXPECT_EQ(orEventCreateWithFlags(&start, orEventEnableTiming), orSuccess);
+
+  EXPECT_EQ(orEventRecord(start, stream), orSuccess);
+
+  // Switch device before creating the end event to force a mismatch.
+  EXPECT_EQ(orSetDevice(1), orSuccess);
+  EXPECT_EQ(orEventCreateWithFlags(&end, orEventEnableTiming), orSuccess);
+  EXPECT_EQ(orSetDevice(0), orSuccess);
+
+  EXPECT_EQ(orEventRecord(end, stream), orSuccess);
+  EXPECT_EQ(orEventSynchronize(start), orSuccess);
+  EXPECT_EQ(orEventSynchronize(end), orSuccess);
+
+  float elapsed_ms = 0.0f;
+  EXPECT_EQ(orEventElapsedTime(&elapsed_ms, start, end), orErrorUnknown);
+
+  EXPECT_EQ(orEventDestroy(start), orSuccess);
+  EXPECT_EQ(orEventDestroy(end), orSuccess);
+  EXPECT_EQ(orStreamDestroy(stream), orSuccess);
+}
+
+TEST_F(EventTest, EventElapsedTimeRequiresTimingFlag) {
+  orStream_t stream = nullptr;
+  EXPECT_EQ(orStreamCreate(&stream), orSuccess);
+
+  orEvent_t start = nullptr;
+  orEvent_t end = nullptr;
+  EXPECT_EQ(orEventCreate(&start), orSuccess);
+  EXPECT_EQ(orEventCreate(&end), orSuccess);
+
+  EXPECT_EQ(orEventRecord(start, stream), orSuccess);
+  EXPECT_EQ(orEventRecord(end, stream), orSuccess);
+  EXPECT_EQ(orEventSynchronize(start), orSuccess);
+  EXPECT_EQ(orEventSynchronize(end), orSuccess);
+
+  // Without timing-enabled events, querying elapsed time must fail.
+  float elapsed_ms = 0.0f;
+  EXPECT_EQ(orEventElapsedTime(&elapsed_ms, start, end), orErrorUnknown);
+
+  EXPECT_EQ(orEventDestroy(start), orSuccess);
+  EXPECT_EQ(orEventDestroy(end), orSuccess);
+  EXPECT_EQ(orStreamDestroy(stream), orSuccess);
+}
+
 TEST_F(EventTest, StreamWaitEvent) {
   orStream_t stream = nullptr;
   EXPECT_EQ(orStreamCreate(&stream), orSuccess);
@@ -85,4 +163,19 @@ TEST_F(EventTest, StreamWaitEvent) {
   EXPECT_EQ(orStreamDestroy(stream), orSuccess);
 }
 
+TEST_F(EventTest, StreamWaitEventInvalidArgs) {
+  orStream_t stream = nullptr;
+  EXPECT_EQ(orStreamCreate(&stream), orSuccess);
+
+  orEvent_t event = nullptr;
+  EXPECT_EQ(orEventCreate(&event), orSuccess);
+
+  // Validate both stream and event inputs for wait calls.
+  EXPECT_EQ(orStreamWaitEvent(nullptr, event, 0), orErrorUnknown);
+  EXPECT_EQ(orStreamWaitEvent(stream, nullptr, 0), orErrorUnknown);
+
+  EXPECT_EQ(orEventDestroy(event), orSuccess);
+  EXPECT_EQ(orStreamDestroy(stream), orSuccess);
+}
+
 } // namespace

test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/tests/memory_tests.cpp

@@ -26,6 +26,12 @@ TEST_F(MemoryTest, AllocateAndFreeHost) {
   EXPECT_EQ(orFreeHost(ptr), orSuccess);
 }
 
+TEST_F(MemoryTest, FreeNullptrIsNoop) {
+  // Freeing a nullptr should behave like CUDA: treated as a no-op success.
+  EXPECT_EQ(orFree(nullptr), orSuccess);
+  EXPECT_EQ(orFreeHost(nullptr), orSuccess);
+}
+
 TEST_F(MemoryTest, AllocateNullptr) {
   EXPECT_EQ(orMalloc(nullptr, 4096), orErrorUnknown);
   EXPECT_EQ(orMallocHost(nullptr, 4096), orErrorUnknown);
@@ -86,6 +92,48 @@ TEST_F(MemoryTest, MemcpyInvalidKind) {
   EXPECT_EQ(orFree(dev_ptr), orSuccess);
 }
 
+TEST_F(MemoryTest, MemcpyInvalidCombinations) {
+  void *dev_src = nullptr, *dev_dst = nullptr;
+  EXPECT_EQ(orMalloc(&dev_src, 8), orSuccess);
+  EXPECT_EQ(orMalloc(&dev_dst, 8), orSuccess);
+
+  char host_buf[8] = {};
+
+  // Deliberately pass mismatched kinds to ensure validation coverage.
+  EXPECT_EQ(
+      orMemcpy(host_buf, dev_src, 4, orMemcpyHostToDevice), orErrorUnknown);
+  EXPECT_EQ(
+      orMemcpy(dev_dst, host_buf, 4, orMemcpyDeviceToHost), orErrorUnknown);
+  EXPECT_EQ(
+      orMemcpy(dev_dst, dev_src, 4, orMemcpyHostToDevice), orErrorUnknown);
+
+  EXPECT_EQ(orFree(dev_src), orSuccess);
+  EXPECT_EQ(orFree(dev_dst), orSuccess);
+}
+
+TEST_F(MemoryTest, MemcpyAsyncHostToDevice) {
+  orStream_t stream = nullptr;
+  EXPECT_EQ(orStreamCreate(&stream), orSuccess);
+
+  const char host_src[] = "async";
+  char host_dst[6] = {};
+  void* dev_ptr = nullptr;
+  EXPECT_EQ(orMalloc(&dev_ptr, sizeof(host_src)), orSuccess);
+
+  // Async copies should complete once the stream is synchronized.
+  EXPECT_EQ(
+      orMemcpyAsync(dev_ptr, host_src, sizeof(host_src), orMemcpyHostToDevice, stream),
+      orSuccess);
+  EXPECT_EQ(orStreamSynchronize(stream), orSuccess);
+  EXPECT_EQ(orMemcpy(
+                host_dst, dev_ptr, sizeof(host_src), orMemcpyDeviceToHost),
+            orSuccess);
+  EXPECT_STREQ(host_dst, host_src);
+
+  EXPECT_EQ(orFree(dev_ptr), orSuccess);
+  EXPECT_EQ(orStreamDestroy(stream), orSuccess);
+}
+
 TEST_F(MemoryTest, PointerAttributes) {
   void* dev_ptr = nullptr;
   EXPECT_EQ(orMalloc(&dev_ptr, 32), orSuccess);
@@ -102,6 +150,14 @@ TEST_F(MemoryTest, PointerAttributes) {
   EXPECT_EQ(orFree(dev_ptr), orSuccess);
 }
 
+TEST_F(MemoryTest, PointerAttributesInvalidArgs) {
+  // Attribute queries must fail on null inputs to avoid dereferencing.
+  char buffer[8] = {};
+  orPointerAttributes attr{};
+  EXPECT_EQ(orPointerGetAttributes(nullptr, buffer), orErrorUnknown);
+  EXPECT_EQ(orPointerGetAttributes(&attr, nullptr), orErrorUnknown);
+}
+
 TEST_F(MemoryTest, ProtectUnprotectDevice) {
   void* dev_ptr = nullptr;
   EXPECT_EQ(orMalloc(&dev_ptr, 64), orSuccess);
@@ -112,4 +168,24 @@ TEST_F(MemoryTest, ProtectUnprotectDevice) {
   EXPECT_EQ(orFree(dev_ptr), orSuccess);
 }
 
+TEST_F(MemoryTest, ProtectReferenceCounting) {
+  void* dev_ptr = nullptr;
+  EXPECT_EQ(orMalloc(&dev_ptr, 64), orSuccess);
+
+  // Call unprotect/protect twice to exercise the refcount transitions.
+  EXPECT_EQ(orMemoryUnprotect(dev_ptr), orSuccess);
+  EXPECT_EQ(orMemoryUnprotect(dev_ptr), orSuccess);
+  EXPECT_EQ(orMemoryProtect(dev_ptr), orSuccess);
+  EXPECT_EQ(orMemoryProtect(dev_ptr), orSuccess);
+
+  EXPECT_EQ(orFree(dev_ptr), orSuccess);
+}
+
+TEST_F(MemoryTest, DoubleFreeFails) {
+  void* dev_ptr = nullptr;
+  EXPECT_EQ(orMalloc(&dev_ptr, 32), orSuccess);
+  EXPECT_EQ(orFree(dev_ptr), orSuccess);
+  EXPECT_EQ(orFree(dev_ptr), orErrorUnknown);
+}
+
 } // namespace

test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/tests/stream_tests.cpp

@@ -21,6 +21,11 @@ TEST_F(StreamTest, StreamCreateAndDestroy) {
   EXPECT_EQ(orStreamDestroy(stream), orSuccess);
 }
 
+TEST_F(StreamTest, StreamCreateNullptr) {
+  // Creation API should reject null double-pointer inputs.
+  EXPECT_EQ(orStreamCreate(nullptr), orErrorUnknown);
+}
+
 TEST_F(StreamTest, StreamCreateWithInvalidPriority) {
   orStream_t stream = nullptr;
   int min_p, max_p;
@@ -30,6 +35,36 @@ TEST_F(StreamTest, StreamCreateWithInvalidPriority) {
   EXPECT_EQ(orStreamCreateWithPriority(&stream, 0, max_p + 1), orErrorUnknown);
 }
 
+TEST_F(StreamTest, StreamCreateWithPriorityValidBounds) {
+  orStream_t stream = nullptr;
+  int min_p, max_p;
+  orDeviceGetStreamPriorityRange(&min_p, &max_p);
+
+  // Lowest priority should be accepted.
+  EXPECT_EQ(orStreamCreateWithPriority(&stream, 0, min_p), orSuccess);
+  EXPECT_EQ(orStreamDestroy(stream), orSuccess);
+
+  // Highest priority should also be accepted.
+  EXPECT_EQ(orStreamCreateWithPriority(&stream, 0, max_p), orSuccess);
+  EXPECT_EQ(orStreamDestroy(stream), orSuccess);
+}
+
+TEST_F(StreamTest, StreamDestroyNullptr) {
+  // Destroying nullptr should follow CUDA error behavior.
+  EXPECT_EQ(orStreamDestroy(nullptr), orErrorUnknown);
+}
+
+TEST_F(StreamTest, StreamGetPriority) {
+  orStream_t stream = nullptr;
+  EXPECT_EQ(orStreamCreate(&stream), orSuccess);
+
+  int priority = -1;
+  EXPECT_EQ(orStreamGetPriority(stream, &priority), orSuccess);
+  EXPECT_EQ(priority, 0);
+
+  EXPECT_EQ(orStreamDestroy(stream), orSuccess);
+}
+
 TEST_F(StreamTest, StreamTaskExecution) {
   orStream_t stream = nullptr;
   EXPECT_EQ(orStreamCreate(&stream), orSuccess);
@@ -43,6 +78,11 @@ TEST_F(StreamTest, StreamTaskExecution) {
   EXPECT_EQ(orStreamDestroy(stream), orSuccess);
 }
 
+TEST_F(StreamTest, AddTaskToStreamNullptr) {
+  // Queueing work should fail fast if the stream handle is invalid.
+  EXPECT_EQ(openreg::addTaskToStream(nullptr, [] {}), orErrorUnknown);
+}
+
 TEST_F(StreamTest, StreamQuery) {
   orStream_t stream = nullptr;
   EXPECT_EQ(orStreamCreate(&stream), orSuccess);
@@ -76,4 +116,18 @@ TEST_F(StreamTest, DeviceSynchronize) {
   EXPECT_EQ(orStreamDestroy(stream2), orSuccess);
 }
 
+TEST_F(StreamTest, DeviceSynchronizeWithNoStreams) {
+  // Even without registered streams, device sync should succeed.
+  EXPECT_EQ(orDeviceSynchronize(), orSuccess);
+}
+
+TEST_F(StreamTest, StreamPriorityRange) {
+  int min_p = -1;
+  int max_p = -1;
+  // OpenReg currently exposes only one priority level; verify the fixed range.
+  EXPECT_EQ(orDeviceGetStreamPriorityRange(&min_p, &max_p), orSuccess);
+  EXPECT_EQ(min_p, 0);
+  EXPECT_EQ(max_p, 0);
+}
+
 } // namespace