mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-20 12:59:49 +00:00
e2753e68bd
* support dynamic tensor descriptor
* use buffer load OOB feature for padding case
* add navi support
* add int8x4 inference kernel
Co-authored-by: Chao Liu <chao@ixt-rack-81.local.lan>
Co-authored-by: Jing Zhang <jizhan@amd.com>
[ROCm/composable_kernel commit: fcbb978828]
242 lines
9.1 KiB
C++
242 lines
9.1 KiB
C++
#ifndef CK_IN_MEMORY_OPERATION_AMD_HPP
|
|
#define CK_IN_MEMORY_OPERATION_AMD_HPP
|
|
|
|
#include "float_type.hpp"
|
|
|
|
#if CK_USE_AMD_BUFFER_ADDRESSING
|
|
#include "amd_buffer_addressing.hpp"
|
|
#include "amd_buffer_addressing_v2.hpp"
|
|
#endif
|
|
|
|
namespace ck {
|
|
|
|
template <typename T>
|
|
__device__ void atomic_add_impl(T* p_dst, T src)
|
|
{
|
|
atomicAdd(p_dst, src);
|
|
}
|
|
|
|
// atomicAdd for float does not support vector type
|
|
template <>
|
|
__device__ void atomic_add_impl<float2_t>(float2_t* p_dst, float2_t src)
|
|
{
|
|
float* p_dst_float = reinterpret_cast<float*>(p_dst);
|
|
const float* p_src_float = reinterpret_cast<const float*>(&src);
|
|
|
|
for(index_t i = 0; i < 2; ++i)
|
|
{
|
|
atomicAdd(&(p_dst_float[i]), p_src_float[i]);
|
|
}
|
|
}
|
|
|
|
template <>
|
|
__device__ void atomic_add_impl<float4_t>(float4_t* p_dst, float4_t src)
|
|
{
|
|
float* p_dst_float = reinterpret_cast<float*>(p_dst);
|
|
const float* p_src_float = reinterpret_cast<const float*>(&src);
|
|
|
|
for(index_t i = 0; i < 4; ++i)
|
|
{
|
|
atomicAdd(&(p_dst_float[i]), p_src_float[i]);
|
|
}
|
|
}
|
|
|
|
template <typename T, index_t DataPerAccess>
|
|
struct SetData
|
|
{
|
|
using vector_t = typename vector_type<T, DataPerAccess>::type;
|
|
|
|
// This version is only for compatibility, don't use this version if possible
|
|
template <AddressSpace SrcAddressSpace, AddressSpace DstAddressSpace>
|
|
__device__ void Run(const T* p_src,
|
|
index_t src_offset,
|
|
bool src_valid,
|
|
index_t /* src_range */,
|
|
T* p_dst,
|
|
index_t dst_offset,
|
|
bool dst_valid,
|
|
index_t /* dst_range */) const
|
|
{
|
|
if(dst_valid)
|
|
{
|
|
if(src_valid)
|
|
{
|
|
#if 0
|
|
*reinterpret_cast<vector_t*>(&p_dst[dst_offset]) =
|
|
*reinterpret_cast<const vector_t*>(&p_src[src_offset]);
|
|
#else
|
|
*reinterpret_cast<vector_t*>(&p_dst[dst_offset]) =
|
|
*reinterpret_cast<const vector_t*>(&p_src[0x3fffffff & src_offset]);
|
|
#endif
|
|
}
|
|
else
|
|
{
|
|
*reinterpret_cast<vector_t*>(&p_dst[dst_offset]) = 0;
|
|
}
|
|
}
|
|
}
|
|
|
|
#if CK_USE_AMD_BUFFER_ADDRESSING
|
|
// buffer_load requires:
|
|
// 1) p_src_thread must be in global memory space, p_dst_thread must be vgpr
|
|
// 2) p_src_thread to be a wavewise pointer.
|
|
// It is user's responsibility to make sure that is true.
|
|
template <>
|
|
__device__ void Run<AddressSpace::Global, AddressSpace::Vgpr>(const T* p_src,
|
|
index_t src_offset,
|
|
bool src_valid,
|
|
index_t src_range,
|
|
T* p_dst,
|
|
index_t dst_offset,
|
|
bool dst_valid,
|
|
index_t /* dst_range */) const
|
|
{
|
|
if(dst_valid)
|
|
{
|
|
*reinterpret_cast<vector_t*>(&p_dst[dst_offset]) =
|
|
amd_buffer_load_v2<T, DataPerAccess>(p_src, src_offset, src_valid, src_range);
|
|
}
|
|
}
|
|
|
|
// buffer_store requires:
|
|
// 1) p_src_thread must be in vgpr space, p_dst_thread must be global memory
|
|
// 2) p_dst_thread to be a wavewise pointer.
|
|
// It is user's responsibility to make sure that is true.
|
|
template <>
|
|
__device__ void Run<AddressSpace::Vgpr, AddressSpace::Global>(const T* p_src,
|
|
index_t src_offset,
|
|
bool src_valid,
|
|
index_t /* src_range */,
|
|
T* p_dst,
|
|
index_t dst_offset,
|
|
bool dst_valid,
|
|
index_t dst_range) const
|
|
{
|
|
const auto zeros = vector_t(0);
|
|
|
|
amd_buffer_store_v2<T, DataPerAccess>(
|
|
src_valid ? *reinterpret_cast<const vector_t*>(&(p_src[src_offset])) : zeros,
|
|
p_dst,
|
|
dst_offset,
|
|
dst_valid,
|
|
dst_range);
|
|
}
|
|
#endif
|
|
};
|
|
|
|
template <typename T, index_t DataPerAccess>
|
|
struct AtomicAddData
|
|
{
|
|
using vector_t = typename vector_type<T, DataPerAccess>::type;
|
|
|
|
// This version is only for compatibility, don't use this version if possible
|
|
template <AddressSpace SrcAddressSpace, AddressSpace DstAddressSpace>
|
|
__device__ void Run(const T* p_src,
|
|
index_t src_offset,
|
|
bool src_valid,
|
|
index_t /* src_range */,
|
|
T* p_dst,
|
|
index_t dst_offset,
|
|
bool dst_valid,
|
|
index_t /* dst_range */) const
|
|
{
|
|
if(src_valid && dst_valid)
|
|
{
|
|
atomic_add_impl(reinterpret_cast<vector_t*>(&p_dst[dst_offset]),
|
|
*reinterpret_cast<const vector_t*>(&p_src[src_offset]));
|
|
}
|
|
}
|
|
|
|
#if CK_USE_AMD_BUFFER_ADDRESSING && CK_USE_AMD_BUFFER_ATOMIC_FADD
|
|
// buffer_atomic requires:
|
|
// 1) p_src_thread must be in vgpr space, p_dst_thread must be global memory
|
|
// 2) p_dst_thread to be a wavewise pointer.
|
|
// It is user's responsibility to make sure that is true.
|
|
template <>
|
|
__device__ void Run<AddressSpace::Vgpr, AddressSpace::Global>(const T* p_src,
|
|
index_t src_offset,
|
|
bool src_valid,
|
|
index_t /* src_range */,
|
|
T* p_dst,
|
|
index_t dst_offset,
|
|
bool dst_valid,
|
|
index_t dst_range) const
|
|
{
|
|
const auto zeros = vector_t(0);
|
|
|
|
amd_buffer_atomic_add<T, DataPerAccess>(
|
|
src_valid ? &(p_src[src_offset]) : &zeros, p_dst, dst_offset, dst_valid, dst_range);
|
|
}
|
|
#endif
|
|
};
|
|
|
|
template <typename T,
|
|
index_t DataPerAccess,
|
|
AddressSpace SrcAddressSpace,
|
|
AddressSpace DstAddressSpace,
|
|
InMemoryDataOperation DstInMemOp,
|
|
index_t SrcDataStride = 1,
|
|
index_t DstDataStride = 1>
|
|
__device__ void transfer_data(const T* p_src,
|
|
index_t src_offset,
|
|
bool src_valid,
|
|
index_t src_range,
|
|
T* p_dst,
|
|
index_t dst_offset,
|
|
bool dst_valid,
|
|
index_t dst_range)
|
|
{
|
|
static_assert(DstInMemOp == InMemoryDataOperation::Set ||
|
|
DstInMemOp == InMemoryDataOperation::AtomicAdd,
|
|
"wrong! InMemoryDataOperation not supported!");
|
|
|
|
// keep it simple, don't use static_if here, otherwise compiler will do weird things
|
|
if constexpr(SrcDataStride == 1 && DstDataStride == 1)
|
|
{
|
|
if constexpr(DstInMemOp == InMemoryDataOperation::Set)
|
|
{
|
|
SetData<T, DataPerAccess>{}.template Run<SrcAddressSpace, DstAddressSpace>(
|
|
p_src, src_offset, src_valid, src_range, p_dst, dst_offset, dst_valid, dst_range);
|
|
}
|
|
else if constexpr(DstInMemOp == InMemoryDataOperation::AtomicAdd)
|
|
{
|
|
AtomicAddData<T, DataPerAccess>{}.template Run<SrcAddressSpace, DstAddressSpace>(
|
|
p_src, src_offset, src_valid, src_range, p_dst, dst_offset, dst_valid, dst_range);
|
|
}
|
|
}
|
|
else
|
|
{
|
|
#pragma unroll
|
|
for(index_t i = 0; i < DataPerAccess; ++i)
|
|
{
|
|
if constexpr(DstInMemOp == InMemoryDataOperation::Set)
|
|
{
|
|
SetData<T, 1>{}.template Run<SrcAddressSpace, DstAddressSpace>(
|
|
p_src,
|
|
src_offset + i * SrcDataStride,
|
|
src_valid,
|
|
src_range,
|
|
p_dst,
|
|
dst_offset + i * DstDataStride,
|
|
dst_valid,
|
|
dst_range);
|
|
}
|
|
else if constexpr(DstInMemOp == InMemoryDataOperation::AtomicAdd)
|
|
{
|
|
AtomicAddData<T, 1>{}.template Run<SrcAddressSpace, DstAddressSpace>(
|
|
p_src,
|
|
src_offset + i * SrcDataStride,
|
|
src_valid,
|
|
src_range,
|
|
p_dst,
|
|
dst_offset + i * DstDataStride,
|
|
dst_valid,
|
|
dst_range);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace ck
|
|
#endif
|