Fix hip usage

nlesc-recruit · Jan 14, 2025 · f23bd0d · f23bd0d
1 parent 7fad49a
commit f23bd0d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/include/cudawrappers/cu.hpp b/include/cudawrappers/cu.hpp
@@ -686,12 +686,18 @@ class GraphKernelNodeParams : public Wrapper<CUDA_KERNEL_NODE_PARAMS> {
                         unsigned blockDimZ, unsigned sharedMemBytes,
                         const std::vector<const void *> &params) {
     _obj.func = function;
+#if defined(__HIP__)
+    _obj.blockDim = {blockDimX, blockDimY, blockDimZ};
+    _obj.gridDim = {gridDimX, gridDimY, gridDimZ};
+
+#else
     _obj.blockDimX = blockDimX;
     _obj.blockDimY = blockDimY;
     _obj.blockDimZ = blockDimZ;
     _obj.gridDimX = gridDimX;
     _obj.gridDimY = gridDimY;
     _obj.gridDimZ = gridDimZ;
+#endif
     _obj.sharedMemBytes = sharedMemBytes;
     _obj.kernelParams = const_cast<void **>(params.data());
     _obj.extra = nullptr;
@@ -763,7 +769,7 @@ class GraphExec : public Wrapper<CUgraphExec> {
 
   explicit GraphExec(const Graph &graph,
                      unsigned int flags = CU_GRAPH_DEFAULT) {
-    checkCudaCall(cuGraphInstantiate(&_obj, graph, flags));
+    checkCudaCall(cuGraphInstantiateWithFlags(&_obj, graph, flags));
   }
 };
 

diff --git a/include/cudawrappers/macros.hpp b/include/cudawrappers/macros.hpp
@@ -661,6 +661,12 @@ typedef uint32_t cuuint32_t;
 #define cuGetErrorName hipDrvGetErrorName
 #define cuGetErrorString hipDrvGetErrorString
 #define cuGetProcAddress hipGetProcAddress
+#define cuGraphAddKernelNode hipGraphAddKernelNode
+#define cuGraphCreate hipGraphCreate
+#define cuGraphDestroy hipGraphDestroy
+#define cuGraphInstantiateWithFlags hipGraphInstantiateWithFlags
+#define cuGraphInstantiate hipGraphInstantiate
+#define cuGraphLaunch hipGraphLaunch
 #define cuInit hipInit
 #define cuIpcCloseMemHandle hipIpcCloseMemHandle
 #define cuIpcGetEventHandle hipIpcGetEventHandle