From 43181845674080fb8a0e1d7ad9649174263ad061 Mon Sep 17 00:00:00 2001
From: Joel Smith <joelsmith@tenstorrent.com>
Date: Sun, 10 Nov 2024 15:09:02 -0800
Subject: [PATCH] Add UNPIN_PAGES tests

---
 test/pin_pages.cpp | 82 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 82 insertions(+)

diff --git a/test/pin_pages.cpp b/test/pin_pages.cpp
index 7739ad5..1c0723a 100644
--- a/test/pin_pages.cpp
+++ b/test/pin_pages.cpp
@@ -357,6 +357,86 @@ void VerifyPinPagesNotContiguous(const EnumeratedDevice &dev)
 
 }
 
+void VerifyUnpinPagesSimple(const EnumeratedDevice &dev)
+{
+    auto page_size = getpagesize();
+
+    struct tenstorrent_pin_pages pin_pages;
+    struct tenstorrent_unpin_pages unpin_pages;
+
+    void *p = std::aligned_alloc(page_size, page_size);
+    std::unique_ptr<void, Freer> page(p);
+
+    DevFd dev_fd(dev.path);
+
+    zero(&pin_pages);
+    pin_pages.in.output_size_bytes = sizeof(pin_pages.out);
+
+    pin_pages.in.flags = TENSTORRENT_PIN_PAGES_CONTIGUOUS;
+    pin_pages.in.virtual_address = reinterpret_cast<uintptr_t>(page.get());
+    pin_pages.in.size = page_size;
+
+    if (ioctl(dev_fd.get(), TENSTORRENT_IOCTL_PIN_PAGES, &pin_pages) != 0)
+        THROW_TEST_FAILURE("PIN_PAGES failed single-page pin.");
+
+    zero(&unpin_pages);
+    unpin_pages.in.virtual_address = reinterpret_cast<uintptr_t>(page.get());
+    unpin_pages.in.size = page_size;
+
+    if (ioctl(dev_fd.get(), TENSTORRENT_IOCTL_UNPIN_PAGES, &unpin_pages) != 0)
+        THROW_TEST_FAILURE("UNPIN_PAGES failed single-page unpin.");
+}
+
+void VerifyUnpinPagesBadSize(const EnumeratedDevice &dev)
+{
+    auto page_size = getpagesize();
+
+    struct tenstorrent_pin_pages pin_pages;
+    struct tenstorrent_unpin_pages unpin_pages;
+
+    void *p = std::aligned_alloc(page_size, page_size);
+    std::unique_ptr<void, Freer> page(p);
+
+    DevFd dev_fd(dev.path);
+
+    zero(&pin_pages);
+    pin_pages.in.output_size_bytes = sizeof(pin_pages.out);
+
+    pin_pages.in.flags = TENSTORRENT_PIN_PAGES_CONTIGUOUS;
+    pin_pages.in.virtual_address = reinterpret_cast<uintptr_t>(page.get());
+    pin_pages.in.size = page_size;
+
+    if (ioctl(dev_fd.get(), TENSTORRENT_IOCTL_PIN_PAGES, &pin_pages) != 0)
+        THROW_TEST_FAILURE("PIN_PAGES failed single-page pin.");
+
+    {
+        zero(&unpin_pages);
+        unpin_pages.in.virtual_address = reinterpret_cast<uintptr_t>(page.get());
+        unpin_pages.in.size = 0;
+
+        if (ioctl(dev_fd.get(), TENSTORRENT_IOCTL_UNPIN_PAGES, &unpin_pages) != -1)
+            THROW_TEST_FAILURE("UNPIN_PAGES succeeded with size = 0");
+    }
+
+    {
+        zero(&unpin_pages);
+        unpin_pages.in.virtual_address = reinterpret_cast<uintptr_t>(page.get());
+        unpin_pages.in.size = page_size / 2;
+
+        if (ioctl(dev_fd.get(), TENSTORRENT_IOCTL_UNPIN_PAGES, &unpin_pages) != -1)
+            THROW_TEST_FAILURE("UNPIN_PAGES succeeded with size = page_size/2.");
+    }
+
+    {
+        zero(&unpin_pages);
+        unpin_pages.in.virtual_address = reinterpret_cast<uintptr_t>(page.get());
+        unpin_pages.in.size = page_size * 2;
+
+        if (ioctl(dev_fd.get(), TENSTORRENT_IOCTL_UNPIN_PAGES, &unpin_pages) != -1)
+            THROW_TEST_FAILURE("UNPIN_PAGES succeeded with size = page_size*2.");
+    }
+}
+
 void TestPinPages(const EnumeratedDevice &dev)
 {
     VerifyPinPagesSimple(dev);
@@ -366,4 +446,6 @@ void TestPinPages(const EnumeratedDevice &dev)
     VerifyPinPagesMultipleRanges(dev);
     VerifyPinPagesContiguous(dev);
     VerifyPinPagesNotContiguous(dev);
+    VerifyUnpinPagesSimple(dev);
+    VerifyUnpinPagesBadSize(dev);
 }