Skip to content

GitLab

"git@developer.sourcefind.cn:orangecat/ollama.git" did not exist on "832b4db9d4baf22497145dde55f334b292ed665f"
Commit 713aa255 authored by Shucai Xiao's avatar Shucai Xiao
Browse files

Merge branch 'argmax_min' of https://github.com/ROCmSoftwarePlatform/AMDMIGraphX into argmax_min

parents e5723db5 d69c8625
Hide whitespace changes
Inline Side-by-side
Showing with 52 additions and 91 deletions
src/include/migraphx/op/argmax.hpp
View file @ 713aa255
......@@ -12,7 +12,7 @@ namespace op {
struct argmax
{
int axis = 0;
int64_t axis = 0;
template <class Self, class F>
static auto reflect(Self& self, F f)
......@@ -25,8 +25,8 @@ struct argmax
shape compute_shape(std::vector<shape> inputs) const
{
check_shapes{inputs, *this}.has(1).standard();
auto lens = inputs[0].lens();
int n_dim = static_cast<int>(lens.size());
auto lens = inputs[0].lens();
int64_t n_dim = static_cast<int64_t>(lens.size());
if(axis >= n_dim || axis < 0)
{
MIGRAPHX_THROW("ARGMAX: axis is out of range.");
......
src/include/migraphx/op/argmin.hpp
View file @ 713aa255
#ifndef MIGRAPHX_GUARD_OPERATORS_ARGMIN_HPP
#define MIGRAPHX_GUARD_OPERATORS_ARGMIN_HPP
//#include <array>
#include <migraphx/operation.hpp>
#include <migraphx/check_shapes.hpp>
//#include <migraphx/stringutils.hpp>
//#include <migraphx/literal.hpp>
#include <migraphx/par_dfor.hpp>
#include <migraphx/config.hpp>
//#include <cmath>
//#include <utility>
namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
......@@ -17,7 +12,7 @@ namespace op {
struct argmin
{
int axis = 0;
int64_t axis = 0;
template <class Self, class F>
static auto reflect(Self& self, F f)
......@@ -30,8 +25,8 @@ struct argmin
shape compute_shape(std::vector<shape> inputs) const
{
check_shapes{inputs, *this}.has(1).standard();
auto lens = inputs[0].lens();
int n_dim = static_cast<int>(lens.size());
auto lens = inputs[0].lens();
int64_t n_dim = static_cast<int64_t>(lens.size());
if(axis >= n_dim || axis < 0)
{
MIGRAPHX_THROW("ARGMIN: axis is out of range.");
......
src/onnx/onnx.cpp
View file @ 713aa255
......@@ -281,10 +281,10 @@ struct onnx_parser
const attribute_map& attributes,
std::vector<instruction_ref> args)
{
int axis = 0;
int64_t axis = 0;
if(contains(attributes, "axis"))
{
axis = parse_value(attributes.at("axis")).at<int>();
axis = static_cast<int64_t>(parse_value(attributes.at("axis")).at<int>());
}
int keep_dims = 1;
......@@ -296,7 +296,7 @@ struct onnx_parser
if(keep_dims == 0)
{
auto ins = prog.add_instruction(op::argmax{axis}, std::move(args));
return prog.add_instruction(op::squeeze{{static_cast<int64_t>(axis)}}, ins);
return prog.add_instruction(op::squeeze{{axis}}, ins);
}
else
{
......@@ -308,10 +308,10 @@ struct onnx_parser
const attribute_map& attributes,
std::vector<instruction_ref> args)
{
int axis = 0;
int64_t axis = 0;
if(contains(attributes, "axis"))
{
axis = parse_value(attributes.at("axis")).at<int>();
axis = static_cast<int64_t>(parse_value(attributes.at("axis")).at<int>());
}
int keep_dims = 1;
......@@ -323,7 +323,7 @@ struct onnx_parser
if(keep_dims == 0)
{
auto ins = prog.add_instruction(op::argmin{axis}, std::move(args));
return prog.add_instruction(op::squeeze{{static_cast<int64_t>(axis)}}, ins);
return prog.add_instruction(op::squeeze{{axis}}, ins);
}
else
{
......
src/targets/cpu/lowering.cpp
View file @ 713aa255
......@@ -650,44 +650,6 @@ struct cpu_logsoftmax
}
};
struct cpu_argmax
{
op::argmax op;
template <class Self, class F>
static auto reflect(Self& self, F f)
{
return migraphx::reflect(self.op, f);
}
std::string name() const { return "cpu::argmax"; }
shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
argument compute(context&, const shape& output_shape, std::vector<argument> args) const
{
return op.compute(output_shape, args);
}
};
struct cpu_argmin
{
op::argmin op;
template <class Self, class F>
static auto reflect(Self& self, F f)
{
return migraphx::reflect(self.op, f);
}
std::string name() const { return "cpu::argmin"; }
shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
argument compute(context&, const shape& output_shape, std::vector<argument> args) const
{
return op.compute(output_shape, args);
}
};
struct cpu_apply
{
program* prog;
......@@ -707,8 +669,6 @@ struct cpu_apply
void init()
{
apply_map["argmax"] = extend_op<cpu_argmax, op::argmax>();
apply_map["argmin"] = extend_op<cpu_argmin, op::argmin>();
apply_map["batch_norm_inference"] =
extend_op<cpu_batch_norm_inference, op::batch_norm_inference>();
apply_map["convolution"] = extend_op<cpu_convolution, op::convolution>();
......
src/targets/gpu/device/argmax.cpp
View file @ 713aa255
......@@ -12,12 +12,9 @@ inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {
namespace device {
void argmax(hipStream_t stream, const argument& result, const argument& arg, int axis)
void argmax(hipStream_t stream, const argument& result, const argument& arg, int64_t axis)
{
arg.visit([&](auto input) {
using type = device_type<std::remove_cv_t<typename decltype(input)::value_type>>;
arg_op<type, argmax_op<type>>(argmax_op<type>{}, stream, result, arg, axis);
});
arg_op(argmax_op{}, stream, result, arg, axis);
}
} // namespace device
......
src/targets/gpu/device/argmin.cpp
View file @ 713aa255
......@@ -12,12 +12,9 @@ inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {
namespace device {
void argmin(hipStream_t stream, const argument& result, const argument& arg, int axis)
void argmin(hipStream_t stream, const argument& result, const argument& arg, int64_t axis)
{
arg.visit([&](auto input) {
using type = device_type<std::remove_cv_t<typename decltype(input)::value_type>>;
arg_op<type, argmin_op<type>>(argmin_op<type>{}, stream, result, arg, axis);
});
arg_op(argmin_op{}, stream, result, arg, axis);
}
} // namespace device
......
src/targets/gpu/include/migraphx/gpu/device/arg_op.hpp
View file @ 713aa255
......@@ -22,8 +22,20 @@ struct val_index
};
template <class T>
MIGRAPHX_DEVICE_CONSTEXPR val_index<T> make_val_index(T v)
{
return {v, -1};
}
template <class T>
MIGRAPHX_DEVICE_CONSTEXPR val_index<T> make_val_index(T v, int64_t i)
{
return {v, i};
}
struct argmax_op
{
template <class T>
MIGRAPHX_DEVICE_CONSTEXPR val_index<T> operator()(val_index<T> x, val_index<T> y) const
{
if(x.val > y.val)
......@@ -36,12 +48,12 @@ struct argmax_op
}
}
MIGRAPHX_DEVICE_CONSTEXPR T init() const { return lowest(); }
MIGRAPHX_DEVICE_CONSTEXPR auto init() const { return lowest(); }
};
template <class T>
struct argmin_op
{
template <class T>
MIGRAPHX_DEVICE_CONSTEXPR val_index<T> operator()(val_index<T> x, val_index<T> y) const
{
if(x.val < y.val)
......@@ -54,11 +66,11 @@ struct argmin_op
}
}
MIGRAPHX_DEVICE_CONSTEXPR T init() const { return highest(); }
MIGRAPHX_DEVICE_CONSTEXPR auto init() const { return highest(); }
};
template <class T, class Op>
void arg_op(Op op, hipStream_t stream, const argument& result, const argument& arg, int axis)
template <class Op>
void arg_op(Op op, hipStream_t stream, const argument& result, const argument& arg, int64_t axis)
{
auto arg_shape = arg.get_shape();
auto lens = arg_shape.lens();
......@@ -69,28 +81,28 @@ void arg_op(Op op, hipStream_t stream, const argument& result, const argument& a
hip_visit_all(arg, arg_shape, batch_shape)([&](auto input, auto arg_s, auto batch_s) {
auto output = device_cast(result.get<int64_t>().data());
using type = device_type<std::remove_cv_t<typename decltype(input)::value_type>>;
// use one block for items in one batch.
const size_t max_block_size = 256;
const std::size_t block_size = compute_block_size(batch_item_num, max_block_size);
gs_launch(stream, batch_shape.elements() * block_size, block_size)(
[=](auto i, auto idx) __device__ {
auto batch_idx = batch_s.multi(i / block_size);
auto data_idx = batch_idx;
T init_val = op.init();
val_index<T> init = {init_val, -1};
gs_launch(stream,
batch_shape.elements() * block_size,
block_size)([=](auto i, auto idx) __device__ {
auto batch_idx = batch_s.multi(i / block_size);
auto data_idx = batch_idx;
auto init = make_val_index<type>(op.init());
auto op_output = block_reduce<max_block_size, Op, val_index<T>>(
idx, op, init, batch_item_num, [&](auto j) __device__ {
data_idx[axis] = j;
T val = input[arg_s.index(data_idx)];
return val_index<T>{val, static_cast<int64_t>(j)};
});
auto op_output =
block_reduce<max_block_size>(idx, op, init, batch_item_num, [&](auto j) __device__ {
data_idx[axis] = j;
return make_val_index(input[arg_s.index(data_idx)], j);
});
if(idx.local == 0)
{
output[batch_s.index(batch_idx)] = op_output.index;
}
});
if(idx.local == 0)
{
output[batch_s.index(batch_idx)] = op_output.index;
}
});
});
}
......
src/targets/gpu/include/migraphx/gpu/device/argmax.hpp
View file @ 713aa255
......@@ -10,7 +10,7 @@ inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {
namespace device {
void argmax(hipStream_t stream, const argument& result, const argument& arg, int axis);
void argmax(hipStream_t stream, const argument& result, const argument& arg, int64_t axis);
} // namespace device
} // namespace gpu
......
src/targets/gpu/include/migraphx/gpu/device/argmin.hpp
View file @ 713aa255
......@@ -10,7 +10,7 @@ inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {
namespace device {
void argmin(hipStream_t stream, const argument& result, const argument& arg, int axis);
void argmin(hipStream_t stream, const argument& result, const argument& arg, int64_t 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