forked from jiuyuan/InfiniTensor
Xpu task4 support: add softmax (#172)
* add softmax on kunlun * format --------- Co-authored-by: Bolun <bolunz@u.nus.edu> Co-authored-by: Haojie Wang <haojie0429@gmail.com>
This commit is contained in:
Bolun Zhang
GitHub
parent
feccd4f318
commit
23b825efc4
|
|
@ -0,0 +1,26 @@
|
|||
#include "operators/softmax.h"
|
||||
#include "kunlun/kunlun_kernel_without_config.h"
|
||||
#include "kunlun/kunlun_runtime.h"
|
||||
|
||||
namespace infini {
|
||||
class SoftmaxXdnn : public KUNLUNKernelWithoutConfig {
|
||||
void compute(const Operator &_op,
|
||||
const RuntimeObj *_context) const override {
|
||||
auto op = as<SoftmaxObj>(_op);
|
||||
auto context = dynamic_cast<const KUNLUNRuntimeObj *>(_context);
|
||||
auto dim = op->getInputs(0)->getDims();
|
||||
auto axis = op->getAxis();
|
||||
|
||||
void *const aData = (op->getInputs(0)->getRawDataPtr<void *>());
|
||||
void *const cData = (op->getOutput()->getRawDataPtr<void *>());
|
||||
|
||||
auto ret = baidu::xpu::api::softmax<float>(
|
||||
context->KUNLUNHandle(), (float *)aData, (float *)cData, dim, axis);
|
||||
assert(ret == 0);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_KERNEL(Device::KUNLUN, OpType::Softmax, DataType::Float32, SoftmaxXdnn,
|
||||
"Softmax_xdnn_KUNLUN_Float32");
|
||||
}; // namespace infini
|
||||
|
|
@ -0,0 +1,136 @@
|
|||
#include "core/graph.h"
|
||||
#include "core/kernel.h"
|
||||
#include "core/runtime.h"
|
||||
#include "kunlun/kunlun_runtime.h"
|
||||
#include "operators/softmax.h"
|
||||
#include "test.h"
|
||||
#include <cmath>
|
||||
namespace infini {
|
||||
|
||||
TEST(XDNN_Softmax, run_axis1) {
|
||||
// Runtime
|
||||
Runtime cpuRuntime = NativeCpuRuntimeObj::getInstance();
|
||||
auto kunlunRuntime = make_ref<KUNLUNRuntimeObj>();
|
||||
|
||||
// Build input data on CPU
|
||||
Tensor inputCpu =
|
||||
make_ref<TensorObj>(Shape{2, 4}, DataType::Float32, cpuRuntime);
|
||||
|
||||
// KUNLUN XPU
|
||||
Graph kunlunGraph = make_ref<GraphObj>(kunlunRuntime);
|
||||
auto inputKunlun = kunlunGraph->cloneTensor(inputCpu);
|
||||
auto kunlunOp = kunlunGraph->addOp<SoftmaxObj>(inputKunlun, nullptr, 1);
|
||||
kunlunGraph->dataMalloc();
|
||||
inputKunlun->copyin(vector<float>{0, 1, 2, 3, 10000, 10001, 10002, 10003});
|
||||
kunlunRuntime->run(kunlunGraph);
|
||||
auto outputKunlun = kunlunOp->getOutput();
|
||||
auto outputKunlun2Cpu = outputKunlun->clone(cpuRuntime);
|
||||
|
||||
// Check
|
||||
EXPECT_TRUE(outputKunlun2Cpu->equalData(
|
||||
vector<float>{0.032058604, 0.08714432, 0.23688284, 0.6439143,
|
||||
0.032058604, 0.08714432, 0.23688284, 0.6439143}));
|
||||
}
|
||||
|
||||
TEST(XDNN_Softmax, run_axis0) {
|
||||
// Runtime
|
||||
Runtime cpuRuntime = NativeCpuRuntimeObj::getInstance();
|
||||
auto kunlunRuntime = make_ref<KUNLUNRuntimeObj>();
|
||||
|
||||
// Build input data on CPU
|
||||
Tensor inputCpu =
|
||||
make_ref<TensorObj>(Shape{2, 4}, DataType::Float32, cpuRuntime);
|
||||
|
||||
// KUNLUN XPU
|
||||
Graph kunlunGraph = make_ref<GraphObj>(kunlunRuntime);
|
||||
auto inputKunlun = kunlunGraph->cloneTensor(inputCpu);
|
||||
auto kunlunOp = kunlunGraph->addOp<SoftmaxObj>(inputKunlun, nullptr, 0);
|
||||
kunlunGraph->dataMalloc();
|
||||
inputKunlun->copyin(vector<float>{0, 1, 2, 3, 10000, 10001, 10002, 10003});
|
||||
kunlunRuntime->run(kunlunGraph);
|
||||
auto outputKunlun = kunlunOp->getOutput();
|
||||
auto outputKunlun2Cpu = outputKunlun->clone(cpuRuntime);
|
||||
|
||||
// Check
|
||||
EXPECT_TRUE(
|
||||
outputKunlun2Cpu->equalData(vector<float>{0., 0., 0., 0., 1, 1, 1, 1}));
|
||||
}
|
||||
|
||||
TEST(XDNN_Softmax2, run_axis1) {
|
||||
// Runtime
|
||||
Runtime cpuRuntime = NativeCpuRuntimeObj::getInstance();
|
||||
auto kunlunRuntime = make_ref<KUNLUNRuntimeObj>();
|
||||
|
||||
// Build input data on CPU
|
||||
Tensor inputCpu =
|
||||
make_ref<TensorObj>(Shape{2, 2, 2, 2}, DataType::Float32, cpuRuntime);
|
||||
|
||||
// KUNLUN XPU
|
||||
Graph kunlunGraph = make_ref<GraphObj>(kunlunRuntime);
|
||||
auto inputKunlun = kunlunGraph->cloneTensor(inputCpu);
|
||||
auto kunlunOp = kunlunGraph->addOp<SoftmaxObj>(inputKunlun, nullptr, 1);
|
||||
kunlunGraph->dataMalloc();
|
||||
inputKunlun->setData(IncrementalGenerator());
|
||||
kunlunRuntime->run(kunlunGraph);
|
||||
auto outputKunlun = kunlunOp->getOutput();
|
||||
auto outputKunlun2Cpu = outputKunlun->clone(cpuRuntime);
|
||||
|
||||
// Check
|
||||
EXPECT_TRUE(outputKunlun2Cpu->equalData(vector<float>{
|
||||
0.0179862, 0.0179862, 0.0179862, 0.0179862, 0.9820138, 0.9820138,
|
||||
0.9820138, 0.9820138, 0.0179862, 0.0179862, 0.0179862, 0.0179862,
|
||||
0.9820138, 0.9820138, 0.9820138, 0.9820138}));
|
||||
}
|
||||
|
||||
TEST(XDNN_Softmax2, run_axis2) {
|
||||
// Runtime
|
||||
Runtime cpuRuntime = NativeCpuRuntimeObj::getInstance();
|
||||
auto kunlunRuntime = make_ref<KUNLUNRuntimeObj>();
|
||||
|
||||
// Build input data on CPU
|
||||
Tensor inputCpu =
|
||||
make_ref<TensorObj>(Shape{2, 2, 2, 2}, DataType::Float32, cpuRuntime);
|
||||
|
||||
// KUNLUN XPU
|
||||
Graph kunlunGraph = make_ref<GraphObj>(kunlunRuntime);
|
||||
auto inputKunlun = kunlunGraph->cloneTensor(inputCpu);
|
||||
auto kunlunOp = kunlunGraph->addOp<SoftmaxObj>(inputKunlun, nullptr, 2);
|
||||
kunlunGraph->dataMalloc();
|
||||
inputKunlun->setData(IncrementalGenerator());
|
||||
kunlunRuntime->run(kunlunGraph);
|
||||
auto outputKunlun = kunlunOp->getOutput();
|
||||
auto outputKunlun2Cpu = outputKunlun->clone(cpuRuntime);
|
||||
|
||||
// Check
|
||||
EXPECT_TRUE(outputKunlun2Cpu->equalData(vector<float>{
|
||||
0.1192029, 0.1192029, 0.8807971, 0.8807971, 0.1192029, 0.1192029,
|
||||
0.8807971, 0.8807971, 0.1192029, 0.1192029, 0.8807971, 0.8807971,
|
||||
0.1192029, 0.1192029, 0.8807971, 0.8807971}));
|
||||
}
|
||||
|
||||
TEST(XDNN_Softmax2, run_axis3) {
|
||||
// Runtime
|
||||
Runtime cpuRuntime = NativeCpuRuntimeObj::getInstance();
|
||||
auto kunlunRuntime = make_ref<KUNLUNRuntimeObj>();
|
||||
|
||||
// Build input data on CPU
|
||||
Tensor inputCpu =
|
||||
make_ref<TensorObj>(Shape{2, 2, 2, 2}, DataType::Float32, cpuRuntime);
|
||||
|
||||
// KUNLUN XPU
|
||||
Graph kunlunGraph = make_ref<GraphObj>(kunlunRuntime);
|
||||
auto inputKunlun = kunlunGraph->cloneTensor(inputCpu);
|
||||
auto kunlunOp = kunlunGraph->addOp<SoftmaxObj>(inputKunlun, nullptr, 3);
|
||||
kunlunGraph->dataMalloc();
|
||||
inputKunlun->setData(IncrementalGenerator());
|
||||
kunlunRuntime->run(kunlunGraph);
|
||||
auto outputKunlun = kunlunOp->getOutput();
|
||||
auto outputKunlun2Cpu = outputKunlun->clone(cpuRuntime);
|
||||
|
||||
// Check
|
||||
EXPECT_TRUE(outputKunlun2Cpu->equalData(vector<float>{
|
||||
0.2689414, 0.7310586, 0.2689414, 0.7310586, 0.2689414, 0.7310586,
|
||||
0.2689414, 0.7310586, 0.2689414, 0.7310586, 0.2689414, 0.7310586,
|
||||
0.2689414, 0.7310586, 0.2689414, 0.7310586}));
|
||||
}
|
||||
} // namespace infini
|
||||
Loading…
Reference in New Issue