Skip to content

GitLab

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

clang format

parent 9bf18316
Hide whitespace changes
Inline Side-by-side
Showing with 36 additions and 42 deletions
src/targets/gpu/device/argmax.cpp
View file @ b222af2f
......@@ -12,10 +12,7 @@ inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {
namespace device {
argument argmax(hipStream_t stream,
const argument& result,
const argument& arg,
int axis)
argument argmax(hipStream_t stream, const argument& result, const argument& arg, int axis)
{
auto lens = arg.get_shape().lens();
auto batch_lens = lens;
......@@ -30,44 +27,47 @@ argument argmax(hipStream_t stream,
hip_tensor_descriptor<n_dim> desc_batch(batch_shape);
hip_tensor_descriptor<n_dim> desc_data(arg.get_shape());
// each block is for one batch
// each block is for one batch
const size_t block_size = 1024;
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>>;
auto batch_idx = desc_batch.multi(blk_idx);
auto data_idx = batch_idx;
auto data_idx = batch_idx;
MIGRAPHX_DEVICE_SHARED type lds_data[block_size];
MIGRAPHX_DEVICE_SHARED int64_t lds_index[block_size];
// load data to lds_data
size_t item_num = n_dims;
for (size_t i = thr_idx; i < n_dims; i += block_size)
for(size_t i = thr_idx; i < n_dims; i += block_size)
{
data_idx[axis] = i;
data_idx[axis] = i;
lds_index[thr_idx] = i;
lds_data[thr_idx] = input_ptr[desc_data.linear(data_idx)];
lds_data[thr_idx] = input_ptr[desc_data.linear(data_idx)];
__syncthreads();
auto size = (item_num > block_size) ? block_size : item_num;
auto size = (item_num > block_size) ? block_size : item_num;
auto stride = (size + 1) / 2;
while (true)
while(true)
{
if (thr_idx + stride < size and lds_data[thr_idx] < lds_data[thr_idx + stride])
if(thr_idx + stride < size and
lds_data[thr_idx] < lds_data[thr_idx + stride])
{
lds_data[thr_idx] = lds_data[thr_idx + stride];
lds_data[thr_idx] = lds_data[thr_idx + stride];
lds_index[thr_idx] = lds_index[thr_idx + stride];
}
__syncthreads();
size = stride;
size = stride;
stride = (stride + 1) / 2;
if (size == 1) break;
if(size == 1)
break;
}
if (thr_idx == 0)
if(thr_idx == 0)
{
output_ptr[blk_idx] = lds_index[0];
}
......
src/targets/gpu/device/argmin.cpp
View file @ b222af2f
......@@ -12,10 +12,7 @@ inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {
namespace device {
argument argmax(hipStream_t stream,
const argument& result,
const argument& arg,
int axis)
argument argmax(hipStream_t stream, const argument& result, const argument& arg, int axis)
{
auto lens = arg.get_shape().lens();
auto batch_lens = lens;
......@@ -30,44 +27,47 @@ argument argmax(hipStream_t stream,
hip_tensor_descriptor<n_dim> desc_batch(batch_shape);
hip_tensor_descriptor<n_dim> desc_data(arg.get_shape());
// each block is for one batch
// each block is for one batch
const size_t block_size = 1024;
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>>;
auto batch_idx = desc_batch.multi(blk_idx);
auto data_idx = batch_idx;
auto data_idx = batch_idx;
MIGRAPHX_DEVICE_SHARED type lds_data[block_size];
MIGRAPHX_DEVICE_SHARED int64_t lds_index[block_size];
// load data to lds_data
size_t item_num = n_dims;
for (size_t i = thr_idx; i < n_dims; i += block_size)
for(size_t i = thr_idx; i < n_dims; i += block_size)
{
data_idx[axis] = i;
data_idx[axis] = i;
lds_index[thr_idx] = i;
lds_data[thr_idx] = input_ptr[desc_data.linear(data_idx)];
lds_data[thr_idx] = input_ptr[desc_data.linear(data_idx)];
__syncthreads();
auto size = (item_num > block_size) ? block_size : item_num;
auto size = (item_num > block_size) ? block_size : item_num;
auto stride = (size + 1) / 2;
while (true)
while(true)
{
if (thr_idx + stride < size and lds_data[thr_idx] > lds_data[thr_idx + stride])
if(thr_idx + stride < size and
lds_data[thr_idx] > lds_data[thr_idx + stride])
{
lds_data[thr_idx] = lds_data[thr_idx + stride];
lds_data[thr_idx] = lds_data[thr_idx + stride];
lds_index[thr_idx] = lds_index[thr_idx + stride];
}
__syncthreads();
size = stride;
size = stride;
stride = (stride + 1) / 2;
if (size == 1) break;
if(size == 1)
break;
}
if (thr_idx == 0)
if(thr_idx == 0)
{
output_ptr[blk_idx] = lds_index[0];
}
......
src/targets/gpu/include/migraphx/gpu/device/argmax.hpp
View file @ b222af2f
......@@ -10,10 +10,7 @@ inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {
namespace device {
argument argmax(hipStream_t stream,
const argument& result,
const argument& arg,
int axis);
argument argmax(hipStream_t stream, const argument& result, const argument& arg, int axis);
} // namespace device
} // namespace gpu
......
src/targets/gpu/include/migraphx/gpu/device/argmin.hpp
View file @ b222af2f
......@@ -10,10 +10,7 @@ inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {
namespace device {
argument argmin(hipStream_t stream,
const argument& result,
const argument& arg,
int axis);
argument argmin(hipStream_t stream, const argument& result, const argument& arg, int axis);
} // namespace device
} // namespace gpu
......
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