Commit bf6fc5f8 authored by Shucai Xiao's avatar Shucai Xiao
Browse files

clang format

parent d1672f1d
......@@ -15,10 +15,10 @@ namespace device {
void argmax(hipStream_t stream, const argument& result, const argument& arg, int axis)
{
auto lens = arg.get_shape().lens();
auto batch_lens = lens;
auto lens = arg.get_shape().lens();
auto batch_lens = lens;
size_t batch_item_num = lens[axis];
batch_lens[axis] = 1;
batch_lens[axis] = 1;
migraphx::shape batch_shape{shape::float_type, batch_lens};
hip_visit_all(result, arg, batch_shape)([&](auto output, auto input, auto batch) {
......@@ -30,24 +30,23 @@ void argmax(hipStream_t stream, const argument& result, const argument& arg, int
block_size *= 2;
}
launch(
stream, batch_shape.elements() * block_size, block_size)([=](auto idx) __device__ {
launch(stream, batch_shape.elements() * block_size, block_size)([=](auto idx) __device__ {
size_t thr_idx = idx.local;
size_t blk_idx = idx.group;
using type = device_type<std::remove_cv_t<typename decltype(output)::value_type>>;
using type = device_type<std::remove_cv_t<typename decltype(output)::value_type>>;
auto batch_idx = batch.multi(blk_idx);
auto data_idx = batch_idx;
MIGRAPHX_DEVICE_SHARED type lds_data[max_block_size + 1];
MIGRAPHX_DEVICE_SHARED int64_t lds_index[max_block_size + 1];
// load data to lds_data
size_t round_item_num = (batch_item_num + block_size - 1) / block_size * block_size;
size_t round_item_num = (batch_item_num + block_size - 1) / block_size * block_size;
size_t remaining_item_num = batch_item_num;
lds_data[max_block_size] = input[0];
lds_data[max_block_size] = input[0];
lds_index[max_block_size] = 0;
for(size_t i = thr_idx; i < round_item_num; i += block_size)
{
if (i < batch_item_num)
if(i < batch_item_num)
{
data_idx[axis] = i;
lds_index[thr_idx] = i;
......@@ -55,13 +54,13 @@ void argmax(hipStream_t stream, const argument& result, const argument& arg, int
}
__syncthreads();
auto item_num = (remaining_item_num > block_size) ? block_size : remaining_item_num;
auto item_num = (remaining_item_num > block_size) ? block_size : remaining_item_num;
reduce_argmax(lds_data, lds_index, block_size, thr_idx, size, max_block_size);
remaining_item_num -= block_size;
}
if (thr_idx == 0)
if(thr_idx == 0)
{
output[batch_idx] = lds_index[max_block_size];
}
......
......@@ -15,10 +15,10 @@ namespace device {
void argmin(hipStream_t stream, const argument& result, const argument& arg, int axis)
{
auto lens = arg.get_shape().lens();
auto batch_lens = lens;
auto lens = arg.get_shape().lens();
auto batch_lens = lens;
size_t batch_item_num = lens[axis];
batch_lens[axis] = 1;
batch_lens[axis] = 1;
migraphx::shape batch_shape{shape::float_type, batch_lens};
hip_visit_all(result, arg, batch_shape)([&](auto output, auto input, auto batch) {
......@@ -30,24 +30,23 @@ void argmin(hipStream_t stream, const argument& result, const argument& arg, int
block_size *= 2;
}
launch(
stream, batch_shape.elements() * block_size, block_size)([=](auto idx) __device__ {
launch(stream, batch_shape.elements() * block_size, block_size)([=](auto idx) __device__ {
size_t thr_idx = idx.local;
size_t blk_idx = idx.group;
using type = device_type<std::remove_cv_t<typename decltype(output)::value_type>>;
using type = device_type<std::remove_cv_t<typename decltype(output)::value_type>>;
auto batch_idx = batch.multi(blk_idx);
auto data_idx = batch_idx;
MIGRAPHX_DEVICE_SHARED type lds_data[max_block_size + 1];
MIGRAPHX_DEVICE_SHARED int64_t lds_index[max_block_size + 1];
// load data to lds_data
size_t round_item_num = (batch_item_num + block_size - 1) / block_size * block_size;
size_t round_item_num = (batch_item_num + block_size - 1) / block_size * block_size;
size_t remaining_item_num = batch_item_num;
lds_data[max_block_size] = input[0];
lds_data[max_block_size] = input[0];
lds_index[max_block_size] = 0;
for(size_t i = thr_idx; i < round_item_num; i += block_size)
{
if (i < batch_item_num)
if(i < batch_item_num)
{
data_idx[axis] = i;
lds_index[thr_idx] = i;
......@@ -55,13 +54,13 @@ void argmin(hipStream_t stream, const argument& result, const argument& arg, int
}
__syncthreads();
auto item_num = (remaining_item_num > block_size) ? block_size : remaining_item_num;
auto item_num = (remaining_item_num > block_size) ? block_size : remaining_item_num;
reduce_argmin(lds_data, lds_index, block_size, thr_idx, size, max_block_size);
remaining_item_num -= block_size;
}
if (thr_idx == 0)
if(thr_idx == 0)
{
output[batch_idx] = lds_index[max_block_size];
}
......
......@@ -38,7 +38,12 @@ inline __device__ void reduce_max(T* data_ptr, size_t block_size, size_t thr_idx
}
template <class T>
inline __device__ void reduce_argmax(T* data_ptr, int64_t* index_ptr, size_t block_size, size_t thr_idx, size_t item_num, size_t max_index)
inline __device__ void reduce_argmax(T* data_ptr,
int64_t* index_ptr,
size_t block_size,
size_t thr_idx,
size_t item_num,
size_t max_index)
{
while(true)
{
......@@ -46,9 +51,9 @@ inline __device__ void reduce_argmax(T* data_ptr, int64_t* index_ptr, size_t blo
auto size = item_num / 2;
for(size_t i = thr_idx; i < size; i += block_size)
{
if (data_ptr[i] < data_ptr[i + stride])
if(data_ptr[i] < data_ptr[i + stride])
{
data_ptr[i] = data_ptr[i + stride];
data_ptr[i] = data_ptr[i + stride];
index_ptr[i] = index_ptr[i + stride];
}
}
......@@ -61,9 +66,9 @@ inline __device__ void reduce_argmax(T* data_ptr, int64_t* index_ptr, size_t blo
if(thr_idx == 0)
{
if (data_ptr[max_index] < data_ptr[0])
if(data_ptr[max_index] < data_ptr[0])
{
data_ptr[max_index] = data_ptr[0];
data_ptr[max_index] = data_ptr[0];
index_ptr[max_index] = index_ptr[0];
}
}
......@@ -72,7 +77,8 @@ inline __device__ void reduce_argmax(T* data_ptr, int64_t* index_ptr, size_t blo
}
template <class T>
inline __device__ void reduce_argmin(T* data_ptr, int64_t* index_ptr, size_t block_size, size_t thr_idx, size_t item_num)
inline __device__ void
reduce_argmin(T* data_ptr, int64_t* index_ptr, size_t block_size, size_t thr_idx, size_t item_num)
{
size_t min_index = item_num;
while(true)
......@@ -81,9 +87,9 @@ inline __device__ void reduce_argmin(T* data_ptr, int64_t* index_ptr, size_t blo
auto size = item_num / 2;
for(size_t i = thr_idx; i < size; i += block_size)
{
if (data_ptr[i] > data_ptr[i + stride])
if(data_ptr[i] > data_ptr[i + stride])
{
data_ptr[i] = data_ptr[i + stride];
data_ptr[i] = data_ptr[i + stride];
index_ptr[i] = index_ptr[i + stride];
}
}
......@@ -96,9 +102,9 @@ inline __device__ void reduce_argmin(T* data_ptr, int64_t* index_ptr, size_t blo
if(thr_idx == 0)
{
if (data_ptr[min_index] > data_ptr[0])
if(data_ptr[min_index] > data_ptr[0])
{
data_ptr[min_index] = data_ptr[0];
data_ptr[min_index] = data_ptr[0];
index_ptr[min_index] = index_ptr[0];
}
}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment