mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-12 11:56:36 +00:00
chore: fotmat code
This commit is contained in:
leejet committed
1 parent
6edeac1b20
commit
5531c8878a
24 files changed
+3052
-3379
No files matched your search
@@ -24,17 +24,17 @@
|
||||
|
||||
namespace sageattention {
|
||||
|
||||
namespace cp_async {
|
||||
namespace cp_async {
|
||||
|
||||
enum class SharedMemFillMode {
|
||||
kFillZero, // Fill zero to shared memory when predicate is false
|
||||
kNoFill // Do not fill zero to shared memory when predicate is false
|
||||
};
|
||||
enum class SharedMemFillMode {
|
||||
kFillZero, // Fill zero to shared memory when predicate is false
|
||||
kNoFill // Do not fill zero to shared memory when predicate is false
|
||||
};
|
||||
|
||||
enum class PrefetchMode {
|
||||
kNoPrefetch, // Do not fetch additional data from global memory to L2
|
||||
kPrefetch // Fetch additional data from global memory to L2
|
||||
};
|
||||
enum class PrefetchMode {
|
||||
kNoPrefetch, // Do not fetch additional data from global memory to L2
|
||||
kPrefetch // Fetch additional data from global memory to L2
|
||||
};
|
||||
|
||||
#if (__CUDACC_VER_MAJOR__ >= 11)
|
||||
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
|
||||
@@ -42,107 +42,107 @@ enum class PrefetchMode {
|
||||
#endif
|
||||
#endif
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted
|
||||
* cp.async instructions to a group
|
||||
*/
|
||||
__device__ __forceinline__ void commit_group() {
|
||||
/*!
|
||||
* \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted
|
||||
* cp.async instructions to a group
|
||||
*/
|
||||
__device__ __forceinline__ void commit_group() {
|
||||
#ifdef CP_ASYNC_ENABLED
|
||||
asm volatile("cp.async.commit_group;\n" ::);
|
||||
asm volatile("cp.async.commit_group;\n" ::);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX cp.async.wait_group instruction
|
||||
* \tparam n Wait till most recent n groups are committed
|
||||
*/
|
||||
template <size_t n>
|
||||
__device__ __forceinline__ void wait_group() {
|
||||
/*!
|
||||
* \brief Wrapper of PTX cp.async.wait_group instruction
|
||||
* \tparam n Wait till most recent n groups are committed
|
||||
*/
|
||||
template <size_t n>
|
||||
__device__ __forceinline__ void wait_group() {
|
||||
#ifdef CP_ASYNC_ENABLED
|
||||
asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
|
||||
asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
|
||||
* global memory to shared memory
|
||||
* \tparam prefetch_mode Whether to fetch additional data from global memory to L2
|
||||
* \tparam T Data type
|
||||
* \param smem_ptr Pointer to shared memory
|
||||
* \param gmem_ptr Pointer to global memory
|
||||
*/
|
||||
template <PrefetchMode prefetch_mode, typename T>
|
||||
__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
|
||||
/*!
|
||||
* \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
|
||||
* global memory to shared memory
|
||||
* \tparam prefetch_mode Whether to fetch additional data from global memory to L2
|
||||
* \tparam T Data type
|
||||
* \param smem_ptr Pointer to shared memory
|
||||
* \param gmem_ptr Pointer to global memory
|
||||
*/
|
||||
template <PrefetchMode prefetch_mode, typename T>
|
||||
__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
|
||||
#ifdef CP_ASYNC_ENABLED
|
||||
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
|
||||
asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
|
||||
"l"(gmem_ptr), "n"(16), "r"(16));
|
||||
} else {
|
||||
asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
|
||||
"l"(gmem_ptr), "n"(16), "r"(16));
|
||||
}
|
||||
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
|
||||
asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
|
||||
"l"(gmem_ptr), "n"(16), "r"(16));
|
||||
} else {
|
||||
asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
|
||||
"l"(gmem_ptr), "n"(16), "r"(16));
|
||||
}
|
||||
#else
|
||||
*((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
|
||||
*((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
|
||||
* global memory to shared memory with predicate.
|
||||
* \tparam prefetch_mode Whether to fetch additional data from global memory to L2
|
||||
* \tparam fill_mode Whether to fill zero to shared memory when predicate is false
|
||||
* \tparam T Data type
|
||||
* \param smem_ptr Pointer to shared memory
|
||||
* \param gmem_ptr Pointer to global memory
|
||||
* \param predicate Predicate value
|
||||
* \note fill zero is slower than not fill zero
|
||||
*/
|
||||
template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
|
||||
__device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) {
|
||||
/*!
|
||||
* \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
|
||||
* global memory to shared memory with predicate.
|
||||
* \tparam prefetch_mode Whether to fetch additional data from global memory to L2
|
||||
* \tparam fill_mode Whether to fill zero to shared memory when predicate is false
|
||||
* \tparam T Data type
|
||||
* \param smem_ptr Pointer to shared memory
|
||||
* \param gmem_ptr Pointer to global memory
|
||||
* \param predicate Predicate value
|
||||
* \note fill zero is slower than not fill zero
|
||||
*/
|
||||
template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
|
||||
__device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) {
|
||||
#ifdef CP_ASYNC_ENABLED
|
||||
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
|
||||
int src_in_bytes = predicate ? 16 : 0;
|
||||
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
|
||||
asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
|
||||
"l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
|
||||
} else {
|
||||
asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
|
||||
"l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
|
||||
}
|
||||
} else {
|
||||
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %0, 0;\n"
|
||||
" @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n"
|
||||
"}\n" ::"r"((int)predicate),
|
||||
"r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
|
||||
} else {
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %0, 0;\n"
|
||||
" @p cp.async.cg.shared.global [%1], [%2], %3;\n"
|
||||
"}\n" ::"r"((int)predicate),
|
||||
"r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
|
||||
}
|
||||
}
|
||||
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
|
||||
int src_in_bytes = predicate ? 16 : 0;
|
||||
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
|
||||
asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
|
||||
"l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
|
||||
} else {
|
||||
asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
|
||||
"l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
|
||||
}
|
||||
} else {
|
||||
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %0, 0;\n"
|
||||
" @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n"
|
||||
"}\n" ::"r"((int)predicate),
|
||||
"r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
|
||||
} else {
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %0, 0;\n"
|
||||
" @p cp.async.cg.shared.global [%1], [%2], %3;\n"
|
||||
"}\n" ::"r"((int)predicate),
|
||||
"r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
|
||||
}
|
||||
}
|
||||
#else
|
||||
if (predicate) {
|
||||
*((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
|
||||
} else {
|
||||
if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
|
||||
*((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0);
|
||||
}
|
||||
}
|
||||
if (predicate) {
|
||||
*((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
|
||||
} else {
|
||||
if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
|
||||
*((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace cp_async
|
||||
} // namespace cp_async
|
||||
|
||||
} // namespace sageattention
|
||||
} // namespace sageattention
|
||||
|
||||
#endif
|
||||
Reference in new issue
Block a user