add rocshmem wait warning

Author: alextmagro

transformer_engine/common/rocshmem_api/rocshmem_waitkernel.hip

@@ -54,6 +54,8 @@ void te_rocshmem_wait_on_stream(uint64_t* sig_addr,
 // ### wait_until_on_stream not yet implemented for rocshmem  ###
 // ### KernelWait is robust but slightly slower due to launch ###
     case WaitKind::ROCSHMEM_WAIT:
+    printf("WARNING: rocshmem wait is not implemented yet, defaulting to
+      kernel wait.\n");
     //   rocshmem__ulonglong_wait_until_on_stream(sig_addr,
     //                                 ROCSHMEM_CMP_EQ,
     //                                 wait_value,

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -304,26 +304,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         "Clean up and finalize the NVSHMEM communication backend and free associated resources",
         py::call_guard<py::gil_scoped_release>());
 #else
-  // rocshmem functions
-  m.def("init_rocshmem_backend", &transformer_engine::pytorch::init_rocshmem_backend,
-        "Initialize ROCSHMEM backend with Pytorch distributed process groups",
-        py::call_guard<py::gil_scoped_release>());
-  m.def("create_rocshmem_tensor", &transformer_engine::pytorch::create_rocshmem_tensor,
-        "Create a tensor in ROCSHMEM shared memory", py::call_guard<py::gil_scoped_release>());
-  m.def("rocshmem_send_on_current_stream",
-        &transformer_engine::pytorch::rocshmem_send_on_current_stream,
-        "Asynchronously send tensor data to a remote PE using ROCSHMEM on the current HIP stream",
-        py::call_guard<py::gil_scoped_release>());
-  m.def("rocshmem_wait_on_current_stream",
-        &transformer_engine::pytorch::rocshmem_wait_on_current_stream,
-        "Wait for a signal value to be updated by a remote PE using ROCSHMEM on the current HIP "
-        "stream",
-        py::call_guard<py::gil_scoped_release>());
-  m.def("rocshmem_finalize", &transformer_engine::pytorch::rocshmem_finalize,
-        "Clean up and finalize the ROCSHMEM communication backend and free associated resources",
-        py::call_guard<py::gil_scoped_release>()); 
-
-  // nvshmem wrappers      
+  // nvshmem/rocshmem wrappers      
   m.def("init_nvshmem_backend", &transformer_engine::pytorch::init_rocshmem_backend,
         "Initialize ROCSHMEM backend with Pytorch distributed process groups",
         py::call_guard<py::gil_scoped_release>());