mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-14 19:18:35 +00:00
tmp code 1
This commit is contained in:
Qun Lin
@@ -123,7 +123,7 @@ using DeviceConvBwdWeightInstance =
|
||||
S<1, 8, 1, 8>,
|
||||
1,
|
||||
ck::BlockGemmPipelineScheduler::Intrawave,
|
||||
ck::BlockGemmPipelineVersion::v2,
|
||||
ck::BlockGemmPipelineVersion::v3,
|
||||
2>;
|
||||
#if 0
|
||||
|
||||
@@ -135,6 +135,465 @@ using DeviceConvBwdWeightInstance =
|
||||
ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, 8 > ;
|
||||
#endif
|
||||
|
||||
namespace ck {
|
||||
namespace tensor_operation {
|
||||
namespace device {
|
||||
|
||||
static constexpr index_t WaveSize = 64;
|
||||
template<typename Argument,
|
||||
typename InDataType>
|
||||
void __device__ load_input_from_global(const Argument* arg, InDataType* p_in, index_t n, uint32_t* p_scratch)
|
||||
{
|
||||
InDataType* p_in_n = p_in + arg->a_g_n_k_wos_strides[1] * n;
|
||||
InDataType* p_in_n_1 = p_in + arg->a_g_n_k_wos_strides[1] * (n + 1);
|
||||
|
||||
const uint32_t W_PACK = 2; //WaveSize / arg->input_spatial_lengths_[1];
|
||||
static_assert(sizeof(InDataType) == 2);
|
||||
|
||||
auto get_offset = [&](index_t y, index_t x)
|
||||
{
|
||||
return y * arg->input_spatial_stride_[0] + x * arg->input_spatial_stride_[1];
|
||||
}
|
||||
for (uint32_t i = 0; i < arg->input_spatial_lengths_[1] / W_PACK; i++)
|
||||
{
|
||||
const index_t offset = get_offset(i * W_PACK + threadIdx.x / (64/W_PACK), threadIdx.x % (64/W_PACK));
|
||||
auto tmp0 = p_in_n[offset];
|
||||
auto tmp1 = p_in_n_1[offset];
|
||||
InDataType* p_scratch_offset = reinterpret_cast<InDataType*>(&p_scratch[i]);
|
||||
p_scratch_offset[0] = tmp1;
|
||||
p_scratch_offset[1] = tmp1;
|
||||
}
|
||||
}
|
||||
|
||||
template<typename Argument,
|
||||
typename OutDataType>
|
||||
void __device__ load_output_from_global(const Argument* arg, OutDataType* p_out, index_t n, uint32_t* p_scratch)
|
||||
{
|
||||
OutDataType* p_out_n = p_out + arg->a_g_n_k_wos_strides[1] * n;
|
||||
OutDataType* p_out_n_1 = p_out + arg->a_g_n_k_wos_strides[1] * (n + 1);
|
||||
|
||||
const uint32_t W_PACK = 2; //WaveSize / arg->input_spatial_lengths_[1];
|
||||
static_assert(sizeof(OutDataType) == 2);
|
||||
|
||||
auto get_offset = [&](index_t y, index_t x)
|
||||
{
|
||||
return y * arg->output_spatial_stride_[0] + x * arg->output_spatial_stride_[1];
|
||||
}
|
||||
for (uint32_t i = 0; i < arg->output_spatial_lengths_[1] / W_PACK; i++)
|
||||
{
|
||||
const index_t offset = get_offset(i * W_PACK + threadIdx.x / (64/W_PACK), threadIdx.x % (64/W_PACK));
|
||||
auto tmp0 = p_out_n[offset];
|
||||
auto tmp1 = p_out_n_1[offset];
|
||||
InDataType* p_scratch_offset = reinterpret_cast<InDataType*>(&p_scratch[i]);
|
||||
p_scratch_offset[0] = tmp1;
|
||||
p_scratch_offset[1] = tmp1;
|
||||
}
|
||||
}
|
||||
|
||||
write_input_to_lds
|
||||
|
||||
template <typename Argument, index_t MinimumOccupancy = 1>
|
||||
__global__ void
|
||||
#if CK_USE_LAUNCH_BOUNDS
|
||||
__launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy)
|
||||
#endif
|
||||
kernel_grouped_conv_bwd_weight_naive(Argument* arg)
|
||||
{
|
||||
#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
|
||||
const index_t g_idx = __builtin_amdgcn_readfirstlane(blockIdx.x);
|
||||
index_t n_idx = 0;
|
||||
|
||||
constexpr index_t Tile_H = 32;
|
||||
constexpr index_t Tile_W = 32;
|
||||
constexpr index_t N_Pack = 2;
|
||||
constexpr index_t SizeOfType = 2;
|
||||
constexpr index_t ShareMemSize = Tile_H * Tile_W * N_Pack * SizeOfType;
|
||||
__shared__ char p_input_0[ShareMemSize];
|
||||
__shared__ char p_input_1[ShareMemSize];
|
||||
__shared__ char p_output_0[ShareMemSize];
|
||||
__shared__ char p_output_1[ShareMemSize];
|
||||
|
||||
constexpr index_t ScratchSize = ShareMemSize / 64 / 4;
|
||||
uint32_t p_input_0_scratch[ScratchSize];
|
||||
uint32_t p_input_1_scratch[ScratchSize];
|
||||
uint32_t p_output_0_scratch[ScratchSize];
|
||||
uint32_t p_output_1_scratch[ScratchSize];
|
||||
|
||||
InDataType* p_in = arg->p_in_grid + g_idx * arg->a_g_n_k_wos_strides[0];
|
||||
OutDataType* p_out = arg->p_out_grid + g_idx * arg->e_g_k_c_xs_strides[0];
|
||||
|
||||
// prefetch 0
|
||||
load_input_from_global(arg, p_in, n_idx, p_input_0_scratch);
|
||||
load_output_from_global(arg, p_out, n_idx, p_output_0_scratch);
|
||||
// prefetch 1
|
||||
load_input_from_global(arg, p_in, n_idx + 1, p_input_1_scratch);
|
||||
load_output_from_global(arg, p_out, n_idx + 1, p_output_0_scratch);
|
||||
|
||||
// write 0
|
||||
write_input_to_lds(arg, p_input_0_scratch);
|
||||
write_output_to_lds(arg, p_output_0_scratch);
|
||||
|
||||
while(num_loop > 0)
|
||||
{
|
||||
// prefetch 0
|
||||
load_input_from_global();
|
||||
load_output_from_global();
|
||||
// do conv_bwd on 0
|
||||
run_conv_bwd_weight();
|
||||
|
||||
// write 1
|
||||
write_input_to_lds();
|
||||
write_output_to_lds();
|
||||
|
||||
// prefetch 1
|
||||
load_input_from_global();
|
||||
load_output_from_global();
|
||||
// do conv_bwd on 1
|
||||
run_conv_bwd_weight();
|
||||
|
||||
// write 0
|
||||
write_input_to_lds();
|
||||
write_output_to_lds();
|
||||
|
||||
num_loop --;
|
||||
};
|
||||
|
||||
if (tail_num == 1)
|
||||
{
|
||||
|
||||
}
|
||||
|
||||
if (tail_num == 2)
|
||||
{
|
||||
|
||||
}
|
||||
|
||||
write_output();
|
||||
|
||||
#else
|
||||
ignore = karg;
|
||||
#endif // end of if (defined(__gfx9__))
|
||||
}
|
||||
|
||||
template <ck::index_t NDimSpatial
|
||||
typename InLayout,
|
||||
typename WeiLayout,
|
||||
typename OutLayout,
|
||||
typename InDataType,
|
||||
typename WeiDataType,
|
||||
typename OutDataType,
|
||||
typename InElementwiseOperation,
|
||||
typename WeiElementwiseOperation,
|
||||
typename OutElementwiseOperation,
|
||||
typename ComputeTypeA = InDataType,
|
||||
typename ComputeTypeB = ComputeTypeA
|
||||
>
|
||||
struct DeviceGroupedConvBwdWeightNaive
|
||||
: public DeviceGroupedConvBwdWeight<NDimSpatial,
|
||||
InLayout,
|
||||
WeiLayout,
|
||||
OutLayout,
|
||||
InDataType,
|
||||
WeiDataType,
|
||||
OutDataType,
|
||||
InElementwiseOperation,
|
||||
WeiElementwiseOperation,
|
||||
OutElementwiseOperation,
|
||||
ComputeTypeA,
|
||||
ComputeTypeB>
|
||||
{
|
||||
static_assert(is_same_v<InElementwiseOperation, element_wise::PassThrough>);
|
||||
static_assert(is_same_v<WeiElementwiseOperation, element_wise::PassThrough>);
|
||||
static_assert(is_same_v<OutElementwiseOperation, element_wise::PassThrough>);
|
||||
|
||||
struct Argument : public BaseArgument
|
||||
{
|
||||
Argument(const InDataType* p_in_grid,
|
||||
WeiDataType* p_wei_grid,
|
||||
const OutDataType* p_out_grid,
|
||||
const std::array<index_t, NDimSpatial + 3>& b_g_n_c_wis_lengths, // input
|
||||
const std::array<index_t, NDimSpatial + 3>& b_g_n_c_wis_strides,
|
||||
const std::array<index_t, NDimSpatial + 3>& e_g_k_c_xs_lengths, // weight
|
||||
const std::array<index_t, NDimSpatial + 3>& e_g_k_c_xs_strides,
|
||||
const std::array<index_t, NDimSpatial + 3>& a_g_n_k_wos_lengths, // output
|
||||
const std::array<index_t, NDimSpatial + 3>& a_g_n_k_wos_strides,
|
||||
const std::array<ck::index_t, NDimSpatial>& conv_filter_strides,
|
||||
const std::array<ck::index_t, NDimSpatial>& conv_filter_dilations,
|
||||
const std::array<ck::index_t, NDimSpatial>& input_left_pads,
|
||||
const std::array<ck::index_t, NDimSpatial>& input_right_pads,
|
||||
InElementwiseOperation in_element_op,
|
||||
WeiElementwiseOperation wei_element_op,
|
||||
OutElementwiseOperation out_element_op,
|
||||
ck::index_t split_k)
|
||||
: p_a_grid_{p_out_grid},
|
||||
p_b_grid_{p_in_grid},
|
||||
p_e_grid_{p_wei_grid},
|
||||
a_grid_desc_k0_m_k1_{},
|
||||
b_grid_desc_k0_n_k1_{},
|
||||
ce_grid_desc_m_n_{},
|
||||
c_grid_desc_mblock_mperblock_nblock_nperblock_{},
|
||||
compute_ptr_offset_of_batch_{},
|
||||
a_element_op_{out_element_op},
|
||||
b_element_op_{in_element_op},
|
||||
cde_element_op_{wei_element_op},
|
||||
Conv_G_{b_g_n_c_wis_lengths[0]},
|
||||
Conv_N_{b_g_n_c_wis_lengths[1]},
|
||||
Conv_K_{e_g_k_c_xs_lengths[1]},
|
||||
Conv_C_{b_g_n_c_wis_lengths[2]},
|
||||
input_spatial_lengths_{},
|
||||
filter_spatial_lengths_{},
|
||||
output_spatial_lengths_{},
|
||||
conv_filter_strides_{conv_filter_strides},
|
||||
input_left_pads_{input_left_pads},
|
||||
input_right_pads_{input_right_pads},
|
||||
k_batch_{split_k}
|
||||
{
|
||||
constexpr index_t spatial_offset = 3;
|
||||
std::copy(begin(b_g_n_c_wis_lengths) + spatial_offset,
|
||||
end(b_g_n_c_wis_lengths),
|
||||
begin(input_spatial_lengths_));
|
||||
std::copy(begin(e_g_k_c_xs_lengths) + spatial_offset,
|
||||
end(e_g_k_c_xs_lengths),
|
||||
begin(filter_spatial_lengths_));
|
||||
std::copy(begin(a_g_n_k_wos_lengths) + spatial_offset,
|
||||
end(a_g_n_k_wos_lengths),
|
||||
begin(output_spatial_lengths_));
|
||||
}
|
||||
|
||||
std::size_t GetWorkspaceSizeBytes() const
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
const ADataType* p_a_grid_;
|
||||
const BDataType* p_b_grid_;
|
||||
EDataType* p_e_grid_;
|
||||
|
||||
index_t M01_;
|
||||
index_t N01_;
|
||||
|
||||
OutElementwiseOperation a_element_op_;
|
||||
InElementwiseOperation b_element_op_;
|
||||
WeiElementwiseOperation cde_element_op_;
|
||||
|
||||
// for checking IsSupportedArgument()
|
||||
const index_t Conv_G_;
|
||||
const index_t Conv_N_;
|
||||
const index_t Conv_K_;
|
||||
const index_t Conv_C_;
|
||||
std::array<ck::index_t, NDimSpatial> input_spatial_lengths_;
|
||||
std::array<ck::index_t, NDimSpatial> filter_spatial_lengths_;
|
||||
std::array<ck::index_t, NDimSpatial> output_spatial_lengths_;
|
||||
const std::array<ck::index_t, NDimSpatial>& conv_filter_strides_;
|
||||
const std::array<ck::index_t, NDimSpatial>& input_left_pads_;
|
||||
const std::array<ck::index_t, NDimSpatial>& input_right_pads_;
|
||||
const index_t k_batch_;
|
||||
};
|
||||
|
||||
// Invoker
|
||||
struct Invoker : public BaseInvoker
|
||||
{
|
||||
using Argument = DeviceOp::Argument;
|
||||
|
||||
void ShowInfo(const Argument& arg)
|
||||
{
|
||||
|
||||
}
|
||||
index_t CalculateGridSize(const Argument& arg)
|
||||
{
|
||||
return arg.Conv_G_;
|
||||
}
|
||||
|
||||
float RunGemmV3(const Argument& arg, const StreamConfig& stream_config = StreamConfig{})
|
||||
{
|
||||
|
||||
index_t gdx = CalculateGridSize(arg);
|
||||
|
||||
float ave_time = 0;
|
||||
|
||||
constexpr index_t minimum_occupancy =
|
||||
BlkGemmPipeSched == BlockGemmPipelineScheduler::Intrawave ? 1 : 2;
|
||||
|
||||
const auto kernel = kernel_grouped_conv_bwd_weight_naive<
|
||||
GridwiseGemm,
|
||||
remove_reference_t<DeviceOp::AGridDesc_K0_M_K1>,
|
||||
remove_reference_t<DeviceOp::BGridDesc_K0_N_K1>,
|
||||
remove_reference_t<
|
||||
DeviceOp::CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock>,
|
||||
ComputePtrOffsetOfStridedBatch<I1, I1, I0>,
|
||||
NumGroupsToMerge,
|
||||
true,
|
||||
InMemoryDataOperationEnum::Set,
|
||||
minimum_occupancy>;
|
||||
|
||||
ave_time += launch_and_time_kernel(
|
||||
stream_config,
|
||||
kernel,
|
||||
dim3(gdx, gdy, gdz),
|
||||
dim3(BlockSize),
|
||||
0,
|
||||
gemm_arg,
|
||||
arg.a_grid_desc_k0_m_k1_,
|
||||
arg.b_grid_desc_k0_n_k1_,
|
||||
arg.c_grid_desc_mblock_mperblock_nblock_nperblock_,
|
||||
arg.compute_ptr_offset_of_batch_,
|
||||
num_k_per_block);
|
||||
|
||||
return ave_time;
|
||||
}
|
||||
|
||||
float Run(const Argument& arg, const StreamConfig& stream_config = StreamConfig{})
|
||||
{
|
||||
float avg_time = 0.f;
|
||||
avg_time += RunGemmV3(arg, stream_config);
|
||||
return avg_time;
|
||||
}
|
||||
|
||||
float Run(const BaseArgument* p_arg,
|
||||
const StreamConfig& stream_config = StreamConfig{}) override
|
||||
{
|
||||
return Run(*dynamic_cast<const Argument*>(p_arg), stream_config);
|
||||
}
|
||||
};
|
||||
|
||||
static constexpr bool IsValidCompilationParameter()
|
||||
{
|
||||
// TODO: properly implement this check
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool IsSupportedArgument(const Argument& arg)
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
bool IsSupportedArgument(const BaseArgument* p_arg) override
|
||||
{
|
||||
return IsSupportedArgument(*dynamic_cast<const Argument*>(p_arg));
|
||||
}
|
||||
|
||||
static auto
|
||||
MakeArgument(const InDataType* p_in_grid,
|
||||
WeiDataType* p_wei_grid,
|
||||
const OutDataType* p_out_grid,
|
||||
const std::array<index_t, NDimSpatial + 3>& b_g_n_c_wis_lengths, // input
|
||||
const std::array<index_t, NDimSpatial + 3>& b_g_n_c_wis_strides,
|
||||
const std::array<index_t, NDimSpatial + 3>& e_g_k_c_xs_lengths, // weight
|
||||
const std::array<index_t, NDimSpatial + 3>& e_g_k_c_xs_strides,
|
||||
const std::array<index_t, NDimSpatial + 3>& a_g_n_k_wos_lengths, // output
|
||||
const std::array<index_t, NDimSpatial + 3>& a_g_n_k_wos_strides,
|
||||
const std::array<ck::index_t, NDimSpatial>& conv_filter_strides,
|
||||
const std::array<ck::index_t, NDimSpatial>& conv_filter_dilations,
|
||||
const std::array<ck::index_t, NDimSpatial>& input_left_pads,
|
||||
const std::array<ck::index_t, NDimSpatial>& input_right_pads,
|
||||
InElementwiseOperation in_element_op,
|
||||
WeiElementwiseOperation wei_element_op,
|
||||
OutElementwiseOperation out_element_op,
|
||||
const ck::index_t split_k)
|
||||
{
|
||||
return Argument{p_in_grid,
|
||||
p_wei_grid,
|
||||
p_out_grid,
|
||||
b_g_n_c_wis_lengths, // input
|
||||
b_g_n_c_wis_strides,
|
||||
e_g_k_c_xs_lengths, // weight
|
||||
e_g_k_c_xs_strides,
|
||||
a_g_n_k_wos_lengths, // output
|
||||
a_g_n_k_wos_strides,
|
||||
conv_filter_strides,
|
||||
conv_filter_dilations,
|
||||
input_left_pads,
|
||||
input_right_pads,
|
||||
in_element_op,
|
||||
wei_element_op,
|
||||
out_element_op,
|
||||
split_k};
|
||||
}
|
||||
|
||||
static auto MakeInvoker() { return Invoker{}; }
|
||||
|
||||
std::unique_ptr<BaseArgument>
|
||||
MakeArgumentPointer(const void* p_in_grid,
|
||||
void* p_wei_grid,
|
||||
const void* p_out_grid,
|
||||
const std::array<index_t, NDimSpatial + 3>& b_g_n_c_wis_lengths, // input
|
||||
const std::array<index_t, NDimSpatial + 3>& b_g_n_c_wis_strides,
|
||||
const std::array<index_t, NDimSpatial + 3>& e_g_k_c_xs_lengths, // weight
|
||||
const std::array<index_t, NDimSpatial + 3>& e_g_k_c_xs_strides,
|
||||
const std::array<index_t, NDimSpatial + 3>& a_g_n_k_wos_lengths, // output
|
||||
const std::array<index_t, NDimSpatial + 3>& a_g_n_k_wos_strides,
|
||||
const std::array<ck::index_t, NDimSpatial>& conv_filter_strides,
|
||||
const std::array<ck::index_t, NDimSpatial>& conv_filter_dilations,
|
||||
const std::array<ck::index_t, NDimSpatial>& input_left_pads,
|
||||
const std::array<ck::index_t, NDimSpatial>& input_right_pads,
|
||||
InElementwiseOperation in_element_op,
|
||||
WeiElementwiseOperation wei_element_op,
|
||||
OutElementwiseOperation out_element_op,
|
||||
const ck::index_t split_k) override
|
||||
{
|
||||
return std::make_unique<Argument>(static_cast<const InDataType*>(p_in_grid),
|
||||
static_cast<WeiDataType*>(p_wei_grid),
|
||||
static_cast<const OutDataType*>(p_out_grid),
|
||||
b_g_n_c_wis_lengths, // input
|
||||
b_g_n_c_wis_strides,
|
||||
e_g_k_c_xs_lengths, // weight
|
||||
e_g_k_c_xs_strides,
|
||||
a_g_n_k_wos_lengths, // output
|
||||
a_g_n_k_wos_strides,
|
||||
conv_filter_strides,
|
||||
conv_filter_dilations,
|
||||
input_left_pads,
|
||||
input_right_pads,
|
||||
in_element_op,
|
||||
wei_element_op,
|
||||
out_element_op,
|
||||
split_k);
|
||||
}
|
||||
|
||||
std::unique_ptr<BaseInvoker> MakeInvokerPointer() override
|
||||
{
|
||||
return std::make_unique<Invoker>(Invoker{});
|
||||
}
|
||||
|
||||
std::string GetTypeString() const override
|
||||
{
|
||||
return "";
|
||||
}
|
||||
|
||||
size_t GetWorkSpaceSize(const BaseArgument* p_arg) const override
|
||||
{
|
||||
auto arg = dynamic_cast<const Argument*>(p_arg);
|
||||
if(arg)
|
||||
{
|
||||
return arg->GetWorkspaceSizeBytes();
|
||||
}
|
||||
else
|
||||
throw std::runtime_error(
|
||||
"The argument pointer is not an object of "
|
||||
"DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle::Argument structure!");
|
||||
}
|
||||
|
||||
void SetWorkSpacePointer(BaseArgument* p_arg,
|
||||
void* p_workspace,
|
||||
const StreamConfig& = StreamConfig{}) const override
|
||||
{
|
||||
auto p_arg_ = dynamic_cast<Argument*>(p_arg);
|
||||
if(p_arg_)
|
||||
{
|
||||
p_arg_->p_workspace_ = p_workspace;
|
||||
}
|
||||
else
|
||||
throw std::runtime_error(
|
||||
"The argument pointer is not an object of "
|
||||
"DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle::Argument structure!");
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <ck::index_t NDimSpatial>
|
||||
using HostConvBwdWeightInstance = ck::tensor_operation::host::ReferenceConvBwdWeight<NDimSpatial,
|
||||
InDataType,
|
||||
|
||||
@@ -156,6 +156,7 @@ function(add_instance_library INSTANCE_NAME)
|
||||
foreach(target IN LISTS INST_TARGETS)
|
||||
string(APPEND offload_targets "--offload-arch=${target} ")
|
||||
endforeach()
|
||||
message(${source}" PROPERTIES COMPILE_FLAGS " ${offload_targets} "supported gpu targets" ${SUPPORTED_GPU_TARGETS})
|
||||
set_source_files_properties(${source} PROPERTIES COMPILE_FLAGS ${offload_targets})
|
||||
list(APPEND INST_OBJ ${source})
|
||||
endforeach()
|
||||
|
||||
Reference in New Issue
Block a user