[DRIVER] Add CUDA P2P support (#209)

2021-08-21 06:00:54 +02:00
parent 6aa5720d75
commit 5b29da719d
7 changed files with 147 additions and 21 deletions
--- a/python/src/triton.cc
+++ b/python/src/triton.cc
@@ -36,6 +36,9 @@ void init_triton_driver(py::module &&m) {
      }))
      .def("max_shared_memory", [](drv::cu_device *self) {
        return self->max_shared_memory();
+      })
+      .def("enable_peer_access", [](drv::cu_device *self, unsigned long long int peer_mem_ptr) {
+        self->enable_peer_access(peer_mem_ptr);
      });
  // host device
  py::class_<drv::host_device, drv::device>(m, "host_device")