[ROCm] add support to ROCm 6.0 and MI300 (#2274)

2024-01-26 15:41:10 -05:00
parent 5265631d15
commit 6b7de1a030
8 changed files with 96 additions and 13 deletions
--- a/csrc/pybind.cpp
+++ b/csrc/pybind.cpp
@@ -81,4 +81,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    "get_device_attribute",
    &get_device_attribute,
    "Gets the specified device attribute.");
+
+  cuda_utils.def(
+    "get_max_shared_memory_per_block_device_attribute",
+    &get_max_shared_memory_per_block_device_attribute,
+    "Gets the maximum shared memory per block device attribute.");
+
 }