[MUSA][8/N] Port CUDA kernels that are compatible with MUSA (#17946)
Signed-off-by: yafeng.li <yafeng.li@mthreads.com> Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
@@ -7,7 +7,7 @@
|
||||
#include <limits>
|
||||
#include <vector>
|
||||
|
||||
#ifndef USE_ROCM
|
||||
#if !defined(USE_ROCM) && !defined(USE_MUSA)
|
||||
#include <dlfcn.h>
|
||||
#define WARP_SIZE 32
|
||||
#include "pytorch_extension_utils.h"
|
||||
@@ -24,7 +24,7 @@ transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_
|
||||
|
||||
#pragma unroll
|
||||
for (int j = lane_id; j < total_chunks; j += WARP_SIZE) {
|
||||
#ifndef USE_ROCM
|
||||
#if !defined(USE_ROCM) && !defined(USE_MUSA)
|
||||
uint64_t tmp;
|
||||
asm volatile("ld.global.nc.b64 %0,[%1];" : "=l"(tmp) : "l"(src + j) : "memory");
|
||||
asm volatile("st.global.cg.b64 [%0],%1;" ::"l"(dst + j), "l"(tmp) : "memory");
|
||||
|
||||
Reference in New Issue
Block a user