[OpenReg] Add tests of device and memory for OpenReg (#161773)

As the title stated.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/161773
Approved by: https://github.com/albanD
ghstack dependencies: #161603, #160099

test/cpp_extensions/open_registration_extension/torch_openreg/third_party/openreg/README.md

@@ -25,7 +25,9 @@ openreg/
 │   └── example.cpp         # Example for OpenReg.
 ├── tests/
 │   ├── event_tests.cpp     # Testcases about OpenReg Event.
-│   └── stream_tests.cpp    # Testcases about OpenReg Stream.
+│   ├── stream_tests.cpp    # Testcases about OpenReg Stream.
+│   ├── device_tests.cpp    # Testcases about OpenReg Device.
+│   └── memory_tests.cpp    # Testcases about OpenReg Memory.
 └── csrc/
     ├── device.cpp          # Implementation of device management APIs
     ├── memory.cpp          # Implementation of memory management APIs

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

@@ -0,0 +1,41 @@
+#include <gtest/gtest.h>
+#include <include/openreg.h>
+
+namespace {
+
+class DeviceTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    orSetDevice(0);
+  }
+};
+
+TEST_F(DeviceTest, GetDeviceCountValid) {
+  int count = -1;
+  EXPECT_EQ(orGetDeviceCount(&count), orSuccess);
+  EXPECT_EQ(count, 2);
+}
+
+TEST_F(DeviceTest, GetDeviceValid) {
+  int device = -1;
+  EXPECT_EQ(orGetDevice(&device), orSuccess);
+  EXPECT_EQ(device, 0);
+}
+
+TEST_F(DeviceTest, SetDeviceValid) {
+  EXPECT_EQ(orSetDevice(1), orSuccess);
+
+  int device = -1;
+  EXPECT_EQ(orGetDevice(&device), orSuccess);
+  EXPECT_EQ(device, 1);
+
+  EXPECT_EQ(orSetDevice(0), orSuccess);
+  EXPECT_EQ(orGetDevice(&device), orSuccess);
+  EXPECT_EQ(device, 0);
+}
+
+TEST_F(DeviceTest, SetDeviceInvalidNegative) {
+  EXPECT_EQ(orSetDevice(-1), orErrorUnknown);
+}
+
+} // namespace

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

@@ -0,0 +1,115 @@
+#include <gtest/gtest.h>
+#include <include/openreg.h>
+
+namespace {
+
+class MemoryTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    orSetDevice(0);
+  }
+};
+
+TEST_F(MemoryTest, AllocateAndFreeDevice) {
+  void* ptr = nullptr;
+  EXPECT_EQ(orMalloc(&ptr, 4096), orSuccess);
+  EXPECT_NE(ptr, nullptr);
+
+  EXPECT_EQ(orFree(ptr), orSuccess);
+}
+
+TEST_F(MemoryTest, AllocateAndFreeHost) {
+  void* ptr = nullptr;
+  EXPECT_EQ(orMallocHost(&ptr, 8192), orSuccess);
+  EXPECT_NE(ptr, nullptr);
+
+  EXPECT_EQ(orFreeHost(ptr), orSuccess);
+}
+
+TEST_F(MemoryTest, AllocateNullptr) {
+  EXPECT_EQ(orMalloc(nullptr, 4096), orErrorUnknown);
+  EXPECT_EQ(orMallocHost(nullptr, 4096), orErrorUnknown);
+}
+
+TEST_F(MemoryTest, AllocateZeroSize) {
+  void* ptr = nullptr;
+  EXPECT_EQ(orMalloc(&ptr, 0), orErrorUnknown);
+  EXPECT_EQ(orMallocHost(&ptr, 0), orErrorUnknown);
+}
+
+TEST_F(MemoryTest, MemcpyHostToDevice) {
+  char host_src[] = "data";
+  char host_dst[5] = {};
+
+  void* dev_ptr = nullptr;
+  EXPECT_EQ(orMalloc(&dev_ptr, 5), orSuccess);
+
+  EXPECT_EQ(orMemcpy(dev_ptr, host_src, 5, orMemcpyHostToDevice), orSuccess);
+  EXPECT_EQ(orMemcpy(host_dst, dev_ptr, 5, orMemcpyDeviceToHost), orSuccess);
+
+  EXPECT_STREQ(host_dst, host_src);
+
+  EXPECT_EQ(orFree(dev_ptr), orSuccess);
+}
+
+TEST_F(MemoryTest, MemcpyDeviceToDevice) {
+  const char host_src[5] = "data";
+  char host_dst[5] = {};
+  void *dev_dst1 = nullptr, *dev_dst2 = nullptr;
+
+  EXPECT_EQ(orMalloc(&dev_dst1, 5), orSuccess);
+  EXPECT_EQ(orMalloc(&dev_dst2, 5), orSuccess);
+
+  EXPECT_EQ(orMemcpy(dev_dst1, host_src, 5, orMemcpyHostToDevice), orSuccess);
+  EXPECT_EQ(orMemcpy(dev_dst2, dev_dst1, 5, orMemcpyDeviceToDevice), orSuccess);
+  EXPECT_EQ(orMemcpy(host_dst, dev_dst2, 5, orMemcpyDeviceToHost), orSuccess);
+
+  EXPECT_STREQ(host_dst, host_src);
+
+  EXPECT_EQ(orFree(dev_dst1), orSuccess);
+  EXPECT_EQ(orFree(dev_dst2), orSuccess);
+}
+
+TEST_F(MemoryTest, MemcpyInvalidKind) {
+  char host_ptr[5] = "data";
+  void* dev_ptr = nullptr;
+
+  EXPECT_EQ(orMalloc(&dev_ptr, 5), orSuccess);
+
+  EXPECT_EQ(
+      orMemcpy(nullptr, host_ptr, 4, orMemcpyHostToDevice), orErrorUnknown);
+  EXPECT_EQ(
+      orMemcpy(dev_ptr, nullptr, 4, orMemcpyHostToDevice), orErrorUnknown);
+  EXPECT_EQ(
+      orMemcpy(dev_ptr, host_ptr, 0, orMemcpyHostToDevice), orErrorUnknown);
+
+  EXPECT_EQ(orFree(dev_ptr), orSuccess);
+}
+
+TEST_F(MemoryTest, PointerAttributes) {
+  void* dev_ptr = nullptr;
+  EXPECT_EQ(orMalloc(&dev_ptr, 32), orSuccess);
+
+  orPointerAttributes attr{};
+  EXPECT_EQ(orPointerGetAttributes(&attr, dev_ptr), orSuccess);
+  EXPECT_EQ(attr.type, orMemoryType::orMemoryTypeDevice);
+  EXPECT_EQ(attr.pointer, dev_ptr);
+
+  char host_ptr[16];
+  EXPECT_EQ(orPointerGetAttributes(&attr, host_ptr), orSuccess);
+  EXPECT_EQ(attr.type, orMemoryType::orMemoryTypeUnmanaged);
+
+  EXPECT_EQ(orFree(dev_ptr), orSuccess);
+}
+
+TEST_F(MemoryTest, ProtectUnprotectDevice) {
+  void* dev_ptr = nullptr;
+  EXPECT_EQ(orMalloc(&dev_ptr, 64), orSuccess);
+
+  EXPECT_EQ(orMemoryUnprotect(dev_ptr), orSuccess);
+  EXPECT_EQ(orMemoryProtect(dev_ptr), orSuccess);
+
+  EXPECT_EQ(orFree(dev_ptr), orSuccess);
+}
+
+} // namespace