forked from jiuyuan/InfiniTensor
83 lines
2.9 KiB
C++
83 lines
2.9 KiB
C++
|
#include "bang/bang_runtime.h"
|
||
|
|
#include "core/graph.h"
|
||
|
|
#include "core/kernel.h"
|
||
|
|
#include "core/runtime.h"
|
||
|
|
#include "operators/reduce.h"
|
||
|
|
|
||
|
|
#include "test.h"
|
||
|
|
|
||
|
|
namespace infini {
|
||
|
|
|
||
|
|
template <typename ReduceObjT>
|
||
|
|
void test_reduce(const Shape &shape, const vector<float> &data,
|
||
|
|
const optional<const vector<int>> &axis, bool keepDims,
|
||
|
|
const vector<float> &ExpectData) {
|
||
|
|
Runtime cpuRuntime = NativeCpuRuntimeObj::getInstance();
|
||
|
|
auto bangRuntime = make_ref<BangRuntimeObj>();
|
||
|
|
|
||
|
|
// Build input data on CPU
|
||
|
|
Tensor icpu = make_ref<TensorObj>(shape, DataType::Float32, cpuRuntime);
|
||
|
|
|
||
|
|
// Build BANG graph
|
||
|
|
Graph g = make_ref<GraphObj>(bangRuntime);
|
||
|
|
auto i = g->cloneTensor(icpu);
|
||
|
|
auto op = g->addOp<ReduceObjT>(i, nullptr, axis, keepDims);
|
||
|
|
|
||
|
|
// allocate BANG memory
|
||
|
|
g->dataMalloc();
|
||
|
|
i->copyin(data);
|
||
|
|
|
||
|
|
// Execute on BANG
|
||
|
|
bangRuntime->run(g);
|
||
|
|
|
||
|
|
// clone BANG output to CPU
|
||
|
|
auto o = op->getOutput();
|
||
|
|
auto ocpu = o->clone(cpuRuntime);
|
||
|
|
|
||
|
|
// check results on CPU
|
||
|
|
EXPECT_TRUE(ocpu->equalData(ExpectData));
|
||
|
|
}
|
||
|
|
|
||
|
|
TEST(BANG_ReduceMean, run) {
|
||
|
|
test_reduce<ReduceMeanObj>(
|
||
|
|
Shape{3, 2, 2}, vector<float>{5, 1, 20, 2, 30, 1, 40, 2, 55, 1, 60, 2},
|
||
|
|
std::nullopt, true, vector<float>{18.25});
|
||
|
|
test_reduce<ReduceMeanObj>(
|
||
|
|
Shape{1, 3, 2, 2, 1},
|
||
|
|
vector<float>{5, 1, 20, 2, 30, 1, 40, 2, 55, 1, 60, 2}, std::nullopt,
|
||
|
|
false, vector<float>{18.25});
|
||
|
|
|
||
|
|
test_reduce<ReduceMeanObj>(
|
||
|
|
Shape{2, 3, 2, 2},
|
||
|
|
vector<float>{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
|
||
|
|
12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23},
|
||
|
|
vector<int>{1, 2}, false, vector<float>{5, 6, 17, 18});
|
||
|
|
test_reduce<ReduceMeanObj>(
|
||
|
|
Shape{2, 3, 2, 2, 1},
|
||
|
|
vector<float>{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
|
||
|
|
12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23},
|
||
|
|
vector<int>{1, 2}, true, vector<float>{5, 6, 17, 18});
|
||
|
|
}
|
||
|
|
|
||
|
|
TEST(BANG_ReduceSum, run) {
|
||
|
|
test_reduce<ReduceSumObj>(Shape{3, 2, 2},
|
||
|
|
vector<float>{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1},
|
||
|
|
std::nullopt, true, vector<float>{12});
|
||
|
|
test_reduce<ReduceSumObj>(Shape{1, 3, 2, 2, 1},
|
||
|
|
vector<float>{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1},
|
||
|
|
std::nullopt, false, vector<float>{12});
|
||
|
|
|
||
|
|
test_reduce<ReduceSumObj>(
|
||
|
|
Shape{2, 3, 2, 2},
|
||
|
|
vector<float>{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
|
||
|
|
12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23},
|
||
|
|
vector<int>{1, 2}, false, vector<float>{30, 36, 102, 108});
|
||
|
|
test_reduce<ReduceSumObj>(
|
||
|
|
Shape{2, 3, 2, 2, 1},
|
||
|
|
vector<float>{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
|
||
|
|
12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23},
|
||
|
|
vector<int>{1, 2}, true, vector<float>{30, 36, 102, 108});
|
||
|
|
}
|
||
|
|
|
||
|
|
} // namespace infini
|