"...targets/git@developer.sourcefind.cn:gaoqiong/migraphx.git" did not exist on "77e80b8ef58e4e087ce9ed6b91953b645cf8dbf0"
Commit ed22ba4f authored by Paul's avatar Paul
Browse files

Format

parent 711aaed9
......@@ -7,51 +7,51 @@
namespace migraphx {
template<class Tensor>
template <class Tensor>
constexpr auto gemm_get_batches()
{
constexpr auto lens = get_shape_c<Tensor>{}.lens;
constexpr auto strides = get_shape_c<Tensor>{}.strides;
constexpr auto new_lens = sequence(lens.size() - _c<2>, [&](auto... is) {
return make_const_array(_c<lens[is]>...);
});
constexpr auto new_strides = sequence(strides.size() - _c<2>, [&](auto... is) {
return make_const_array(_c<strides[is]>...);
});
constexpr auto lens = get_shape_c<Tensor>{}.lens;
constexpr auto strides = get_shape_c<Tensor>{}.strides;
constexpr auto new_lens = sequence(
lens.size() - _c<2>, [&](auto... is) { return make_const_array(_c<lens[is]>...); });
constexpr auto new_strides = sequence(
strides.size() - _c<2>, [&](auto... is) { return make_const_array(_c<strides[is]>...); });
return make_shape(new_lens, new_strides);
}
template<class Tensor>
template <class Tensor>
constexpr auto gemm_get_matrix()
{
constexpr auto lens = get_shape_c<Tensor>{}.lens;
constexpr auto strides = get_shape_c<Tensor>{}.strides;
constexpr auto m = lens.size() - _c<2>;
constexpr auto n = lens.size() - _c<1>;
constexpr auto new_lens = make_const_array(_c<lens[m]>, _c<lens[n]>);
constexpr auto lens = get_shape_c<Tensor>{}.lens;
constexpr auto strides = get_shape_c<Tensor>{}.strides;
constexpr auto m = lens.size() - _c<2>;
constexpr auto n = lens.size() - _c<1>;
constexpr auto new_lens = make_const_array(_c<lens[m]>, _c<lens[n]>);
constexpr auto new_strides = make_const_array(_c<strides[m]>, _c<strides[n]>);
return make_shape(new_lens, new_strides);
}
template<class Tensor, class T>
template <class Tensor, class T>
constexpr auto gemm_batch_slice(Tensor t, T i)
{
constexpr auto batch = gemm_get_batches<Tensor>();
constexpr auto batch = gemm_get_batches<Tensor>();
constexpr auto matrix = gemm_get_matrix<Tensor>();
return make_tensor_view(t.data() + batch.index(i), matrix);
}
template<class BlocksPerBatch, class T, class... Ts>
template <class BlocksPerBatch, class T, class... Ts>
constexpr auto gemm_batch_args(index idx, BlocksPerBatch bpb, T x, Ts... xs)
{
return [=](auto f) {
// All tensors should have the same rank
static_assert((true and ... and (get_shape_c<T>{}.lens.size() == get_shape_c<Ts>{}.lens.size())));
static_assert(
(true and ... and (get_shape_c<T>{}.lens.size() == get_shape_c<Ts>{}.lens.size())));
if constexpr(get_shape_c<T>{}.lens.size() > 2)
{
// Get the first batch since all batches should have the same number of elements
constexpr auto batch = gemm_get_batches<T>();
static_assert((true and ... and (batch.elements() == gemm_get_batches<Ts>().elements())));
static_assert(
(true and ... and (batch.elements() == gemm_get_batches<Ts>().elements())));
idx.group_stride(bpb * batch.elements(), [&](auto gidx) {
const auto batch_idx = gidx / bpb;
f(gemm_batch_slice(x, batch_idx), gemm_batch_slice(xs, batch_idx)...);
......@@ -59,7 +59,7 @@ constexpr auto gemm_batch_args(index idx, BlocksPerBatch bpb, T x, Ts... xs)
}
else
{
f(x, xs...);
f(x, xs...);
}
};
}
......
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