Add cross-platform support for Marlin kernel attribute setting

Modify the Marlin kernel to support both CUDA and ROCm platforms by:
- Introducing platform-specific function attribute setting
- Using conditional compilation with #ifdef __HIP_PLATFORM_AMD__
- Refactoring kernel attribute configuration for better cross-platform compatibility

Author: petrex

torchao/csrc/cuda/sparse_marlin/marlin_kernel_nm.cu

@@ -861,10 +861,28 @@ __global__ void Marlin_24(
            thread_n_blocks == THREAD_N_BLOCKS &&                              \
            thread_k_blocks == THREAD_K_BLOCKS &&                              \
            group_blocks == GROUP_BLOCKS) {                                    \
-    cudaFuncSetAttribute(                                                     \
-        Marlin_24<NUM_BITS, THREADS, THREAD_N_BLOCKS, THREAD_M_BLOCKS,        \
-                  THREAD_K_BLOCKS, STAGES, GROUP_BLOCKS>,                     \
-        cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem);         \
+    /* Define the kernel type */                                              \
+    using KernelType = decltype(Marlin_24<NUM_BITS, THREADS, THREAD_N_BLOCKS, \
+                                         THREAD_M_BLOCKS, THREAD_K_BLOCKS,    \
+                                         STAGES, GROUP_BLOCKS>);              \
+    /* Get function pointer for the kernel */                                 \
+    const void* kernel_ptr = reinterpret_cast<const void*>(                   \
+        &Marlin_24<NUM_BITS, THREADS, THREAD_N_BLOCKS, THREAD_M_BLOCKS,       \
+                  THREAD_K_BLOCKS, STAGES, GROUP_BLOCKS>);                    \
+                                                                              \
+    /* Set attribute based on platform */                                     \
+    #ifdef __HIP_PLATFORM_AMD__                                               \
+      hipFuncSetAttribute(kernel_ptr,                                         \
+                         hipFuncAttributeMaxDynamicSharedMemorySize,          \
+                         max_shared_mem);                                     \
+    #else                                                                     \
+      cudaFuncSetAttribute(                                                   \
+          Marlin_24<NUM_BITS, THREADS, THREAD_N_BLOCKS, THREAD_M_BLOCKS,      \
+                    THREAD_K_BLOCKS, STAGES, GROUP_BLOCKS>,                   \
+          cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem);       \
+    #endif                                                                    \
+                                                                              \
+    /* Launch kernel */                                                       \
     Marlin_24<NUM_BITS, THREADS, THREAD_N_BLOCKS, THREAD_M_BLOCKS,            \
               THREAD_K_BLOCKS, STAGES, GROUP_BLOCKS>                          \
         <<<blocks, THREADS, max_shared_mem, stream>>>(A_ptr, B_ptr, meta_ptr, \