From 9db6703b585ba77df9b15c9d527f426da23c88a3 Mon Sep 17 00:00:00 2001
From: OdinaryWord <sx-hz@163.com>
Date: Mon, 29 Jan 2024 15:07:49 +0800
Subject: [PATCH] add reshape

---
 examples/NNmodel                           |  2 +-
 src/kernels/ascend/reshape.cc              | 71 ++++++++++++++++++++++
 test/kernels/ascend/test_ascend_softmax.cc |  6 ++
 3 files changed, 78 insertions(+), 1 deletion(-)
 create mode 100644 src/kernels/ascend/reshape.cc

diff --git a/examples/NNmodel b/examples/NNmodel
index 51d31052..b896cec2 160000
--- a/examples/NNmodel
+++ b/examples/NNmodel
@@ -1 +1 @@
-Subproject commit 51d3105277f3774ed31c02ed4cd11fa92925af77
+Subproject commit b896cec2dba5b8522b141ac4f89eb43074ee1b98
diff --git a/src/kernels/ascend/reshape.cc b/src/kernels/ascend/reshape.cc
new file mode 100644
index 00000000..ff72c250
--- /dev/null
+++ b/src/kernels/ascend/reshape.cc
@@ -0,0 +1,71 @@
+#include "operators/reshape.h"
+#include "aclnnop/level2/aclnn_copy.h"
+#include "ascend/ascend_kernel_without_config.h"
+#include "ascend/ascend_runtime.h"
+
+namespace infini {
+
+class CopyAclnn : public ASCENDKernelWithoutConfig {
+
+    void compute(const Operator &_op,
+                 const RuntimeObj *_context) const override {
+        auto op = as<MatmulObj>(_op);
+        auto context = dynamic_cast<const ASCENDRuntimeObj *>(_context);
+
+        void *const aData = (op->getInputs(0)->getRawDataPtr<void *>());
+        void *const bData = (op->getInputs(1)->getRawDataPtr<void *>());
+        void *const cData = (op->getOutput()->getRawDataPtr<void *>());
+
+        auto selfD = op->getInputs(0)->getDims();
+        auto selfS = op->getInputs(0)->getStride();
+        auto matD = op->getInputs(1)->getDims();
+        auto matS = op->getInputs(1)->getStride();
+        auto outD = op->getOutput()->getDims();
+        auto outS = op->getOutput()->getStride();
+
+        std::vector<int64_t> selfDim = MycastTo64(selfD);
+        std::vector<int64_t> selfStride = MycastTo64(selfS);
+        std::vector<int64_t> matDim = MycastTo64(matD);
+        std::vector<int64_t> matStride = MycastTo64(matS);
+        std::vector<int64_t> outputDim = MycastTo64(outD);
+        std::vector<int64_t> outputStride = MycastTo64(outS);
+
+        auto selfTensor = aclCreateTensor(
+            selfDim.data(), selfDim.size(), ACL_FLOAT, selfStride.data(), 0,
+            aclFormat::ACL_FORMAT_ND, selfDim.data(), selfDim.size(), aData);
+        auto matTensor = aclCreateTensor(
+            matDim.data(), matDim.size(), ACL_FLOAT, matStride.data(), 0,
+            aclFormat::ACL_FORMAT_ND, matDim.data(), matDim.size(), bData);
+        auto outputTensor =
+            aclCreateTensor(outputDim.data(), outputDim.size(), ACL_FLOAT,
+                            outputStride.data(), 0, aclFormat::ACL_FORMAT_ND,
+                            outputDim.data(), outputDim.size(), cData);
+
+        uint64_t workspaceSize = 0;
+        aclOpExecutor *executor;
+
+        auto ret = aclnnMatmulGetWorkspaceSize(
+            selfTensor, matTensor, outputTensor, 1, &workspaceSize, &executor);
+        void *workspaceAddr = nullptr;
+        if (workspaceSize > 0) {
+            workspaceAddr = context->getWorkspace(workspaceSize);
+        }
+        assert(ret == ACL_SUCCESS);
+        ret = aclnnMatmul(workspaceAddr, workspaceSize, executor,
+                          context->ASCENDHandle());
+        assert(ret == ACL_SUCCESS);
+
+        ret = aclrtSynchronizeStream(context->ASCENDHandle());
+        assert(ret == ACL_SUCCESS);
+
+        // aclDestroyTensor(selfTensor);
+        // aclDestroyTensor(matTensor);
+        // aclDestroyTensor(outputTensor);
+
+        return;
+    }
+};
+
+REGISTER_KERNEL(Device::ASCEND, OpType::MatMul, MatmulAclnn,
+                "matmul_ASCEND_float");
+}; // namespace infini
diff --git a/test/kernels/ascend/test_ascend_softmax.cc b/test/kernels/ascend/test_ascend_softmax.cc
index d8371dd4..40b2f611 100644
--- a/test/kernels/ascend/test_ascend_softmax.cc
+++ b/test/kernels/ascend/test_ascend_softmax.cc
@@ -43,6 +43,12 @@ TEST(ascend_ElementWise, run) {
                       0.9820138, 0.9820138, 0.9820138, 0.0179862, 0.0179862,
                       0.0179862, 0.0179862, 0.9820138, 0.9820138, 0.9820138,
                       0.9820138});
+    testSoftmax<SoftmaxObj>(
+        IncrementalGenerator(), Shape{2, 2, 2, 2}, 2,
+        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});
     testSoftmax<SoftmaxObj>(
         IncrementalGenerator(), Shape{2, 2, 2, 2}, 3,
         vector<float>{0.2689414, 0.7310586, 0.2689414, 0.7310586, 0.2689414,