Merge pull request #1782 from IntelPython/add-usm-memory-constructor-from-existing-allocation

Add usm memory constructor from existing allocation

Author: oleksandr-pavlyk

dpctl/apis/include/dpctl4pybind11.hpp

@@ -27,6 +27,7 @@
 
 #include "dpctl_capi.h"
 #include <complex>
+#include <exception>
 #include <memory>
 #include <pybind11/pybind11.h>
 #include <sycl/sycl.hpp>
@@ -748,6 +749,53 @@ class usm_memory : public py::object
             throw py::error_already_set();
     }
 
+    /*! @brief Create usm_memory object from shared pointer that manages
+     *  lifetime of the USM allocation.
+     */
+    usm_memory(void *usm_ptr,
+               size_t nbytes,
+               const sycl::queue &q,
+               std::shared_ptr<void> shptr)
+    {
+        auto const &api = ::dpctl::detail::dpctl_capi::get();
+        DPCTLSyclUSMRef usm_ref = reinterpret_cast<DPCTLSyclUSMRef>(usm_ptr);
+        auto q_uptr = std::make_unique<sycl::queue>(q);
+        DPCTLSyclQueueRef QRef =
+            reinterpret_cast<DPCTLSyclQueueRef>(q_uptr.get());
+
+        auto vacuous_destructor = []() {};
+        py::capsule mock_owner(vacuous_destructor);
+
+        // create memory object owned by mock_owner, it is a new reference
+        PyObject *_memory =
+            api.Memory_Make_(usm_ref, nbytes, QRef, mock_owner.ptr());
+        auto ref_count_decrementer = [](PyObject *o) noexcept { Py_DECREF(o); };
+
+        using py_uptrT =
+            std::unique_ptr<PyObject, decltype(ref_count_decrementer)>;
+
+        if (!_memory) {
+            throw py::error_already_set();
+        }
+
+        auto memory_uptr = py_uptrT(_memory, ref_count_decrementer);
+        std::shared_ptr<void> *opaque_ptr = new std::shared_ptr<void>(shptr);
+
+        Py_MemoryObject *memobj = reinterpret_cast<Py_MemoryObject *>(_memory);
+        // replace mock_owner capsule as the owner
+        memobj->refobj = Py_None;
+        // set opaque ptr field, usm_memory now knowns that USM is managed
+        // by smart pointer
+        memobj->_opaque_ptr = reinterpret_cast<void *>(opaque_ptr);
+
+        // _memory will delete created copies of sycl::queue, and
+        // std::shared_ptr and the deleter of the shared_ptr<void> is
+        // supposed to free the USM allocation
+        m_ptr = _memory;
+        q_uptr.release();
+        memory_uptr.release();
+    }
+
     sycl::queue get_queue() const
     {
         Py_MemoryObject *mem_obj = reinterpret_cast<Py_MemoryObject *>(m_ptr);

examples/pybind11/external_usm_allocation/README.md

@@ -1,8 +1,10 @@
 # Exposing USM Allocations Made by the Native Code to dpctl
 
-This extension demonstrates how a Python object backed by
+This extension demonstrates how a Python object representing
 a native class, which allocates USM memory, can expose it
-to the `dpctl.memory` entities using `__sycl_usm_array_interface__`.
+to the `dpctl.memory` entities using `__sycl_usm_array_interface__`,
+and how to create `dpctl.memory` object from allocation made
+in native extension.
 
 
 ## Building
@@ -29,4 +31,10 @@ shared
 [1.0, 1.0, 0.0, 2.0, 2.0]
 [0.0, 0.0, 0.0, 3.0, -1.0]
 [0.0, 0.0, 0.0, -1.0, 5.0]
+
+========================================
+device
+64
+[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
+ 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
 ```

examples/pybind11/external_usm_allocation/example.py

@@ -50,3 +50,10 @@
 list_of_lists = matr.tolist()
 for row in list_of_lists:
     print(row)
+
+print("====" * 10)
+
+mbuf = eua.make_zeroed_device_memory(4 * 16, q)
+print(mbuf.get_usm_type())
+print(mbuf.nbytes)
+print(mbuf.copy_to_host())

examples/pybind11/external_usm_allocation/external_usm_allocation/__init__.py

@@ -16,9 +16,9 @@
 
 # coding: utf-8
 
-from ._external_usm_alloc import DMatrix
+from ._external_usm_alloc import DMatrix, make_zeroed_device_memory
 
-__all__ = ["DMatrix"]
+__all__ = ["DMatrix", "make_zeroed_device_memory"]
 
 __doc__ = """
    Example of implementing C++ class with its own USM memory allocation logic

examples/pybind11/external_usm_allocation/external_usm_allocation/_usm_alloc_example.cpp

@@ -31,11 +31,12 @@
 /// content of the object into list of lists of Python floats.
 ///
 //===----------------------------------------------------------------------===//
-#include <CL/sycl.hpp>
+#include <sycl/sycl.hpp>
 
 #include "dpctl4pybind11.hpp"
 #include "pybind11/pybind11.h"
 #include "pybind11/stl.h"
+#include <memory>
 
 namespace py = pybind11;
 
@@ -120,6 +121,26 @@ py::list tolist(DMatrix &m)
     return rows;
 }
 
+dpctl::memory::usm_memory make_zeroed_device_memory(size_t nbytes,
+                                                    sycl::queue &q)
+{
+    char *data = sycl::malloc_device<char>(nbytes, q);
+    q.memset(data, 0, nbytes).wait();
+
+    const sycl::context &ctx = q.get_context();
+    auto _deleter = [ctx](void *ptr) {
+        try {
+            ::sycl::free(ptr, ctx);
+        } catch (const std::exception &e) {
+            std::cout << "Call to sycl::free caught an exception: " << e.what()
+                      << std::endl;
+        }
+    };
+    auto shptr = std::shared_ptr<void>(data, _deleter);
+
+    return dpctl::memory::usm_memory(data, nbytes, q, shptr);
+}
+
 PYBIND11_MODULE(_external_usm_alloc, m)
 {
     py::class_<DMatrix> dm(m, "DMatrix");
@@ -128,4 +149,7 @@ PYBIND11_MODULE(_external_usm_alloc, m)
     dm.def_property("__sycl_usm_array_interface__", &construct_sua_iface,
                     nullptr);
     dm.def("tolist", &tolist, "Return matrix a Python list of lists");
+
+    m.def("make_zeroed_device_memory", &make_zeroed_device_memory,
+          "Returns zero-initialized USM-device allocation created C++");
 }

examples/pybind11/external_usm_allocation/tests/test_direct.py

@@ -0,0 +1,37 @@
+#                      Data Parallel Control (dpctl)
+#
+# Copyright 2020-2024 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# coding: utf-8
+
+import external_usm_allocation as eua
+
+import dpctl
+import dpctl.memory as dpm
+import dpctl.tensor as dpt
+
+
+def test_direct():
+    q = dpctl.SyclQueue()
+
+    nb = 2 * 30
+    mbuf = eua.make_zeroed_device_memory(nb, q)
+
+    assert isinstance(mbuf, dpm.MemoryUSMDevice)
+    assert mbuf.nbytes == 2 * 30
+    assert mbuf.sycl_queue == q
+
+    x = dpt.usm_ndarray(30, dtype="i2", buffer=mbuf)
+    assert dpt.all(x == dpt.zeros(30, dtype="i2", sycl_queue=q))