zzf/add op tril (#40)

* add op tril --------- Co-authored-by: wugeshui <[email protected]>
DeepLink-org · May 9, 2023 · b702b11 · b702b11
1 parent c592042
commit b702b11
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 15 deletions.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -212,7 +212,8 @@ jobs:
           set -e
           export CI=true
           cd ${NFS_PATH}/${GITHUB_RUN_NUMBER} && ls
-          rsync -a ${CLUSTER_1984}:${NFS_PATH}/${GITHUB_RUN_NUMBER}/${BUILD_TEST1}/${GEN_DATA} ${NFS_PATH}/${GITHUB_RUN_NUMBER}/${BUILD_TEST1}/DIOPI-TEST/python/
+          echo "rsync -avz ${CLUSTER_1984}:${NFS_PATH}/${GITHUB_RUN_NUMBER}/${BUILD_TEST1}/${GEN_DATA} ${NFS_PATH}/${GITHUB_RUN_NUMBER}/${BUILD_TEST1}/DIOPI-TEST/python/"
+          rsync -avz ${CLUSTER_1984}:${NFS_PATH}/${GITHUB_RUN_NUMBER}/${BUILD_TEST1}/${GEN_DATA} ${NFS_PATH}/${GITHUB_RUN_NUMBER}/${BUILD_TEST1}/DIOPI-TEST/python/
           source /mnt/cache/share/platform/env/camb_ci_diopi_impl && cd ${NFS_PATH}/${GITHUB_RUN_NUMBER} && cd ${BUILD_TEST1}
           srun --job-name=${GITHUB_JOB} --partition=${SLURM_PAR_CAMB} --time=10 --gres=mlu:${GPU_REQUESTS} bash -c 'cd DIOPI-TEST/python && python main.py --mode run_test --impl_folder ${NFS_PATH_LUSTRE}/${GITHUB_RUN_NUMBER}/${BUILD_TEST1}/DIOPI-IMPL/camb/' \
           && cd ${NFS_PATH}/${GITHUB_RUN_NUMBER}/${BUILD_TEST1} && git clean -xdf ${GEN_DATA} || ( cd ${NFS_PATH}/${GITHUB_RUN_NUMBER}/${BUILD_TEST1} && git clean -xdf ${GEN_DATA} && exit 1 )

diff --git a/DIOPI-IMPL/camb/device_configs.py b/DIOPI-IMPL/camb/device_configs.py
@@ -504,20 +504,6 @@
         ),
     ),
 
-    'tril': dict(
-        name=["tril"],
-        tensor_para=dict(
-            args=[
-                {
-                    "ins": ['input'],
-                    "dtype": [Skip(Dtype.float64), Skip(Dtype.float32), Skip(Dtype.float16),
-                              Skip(Dtype.int64), Skip(Dtype.int32), Skip(Dtype.int16),
-                              Skip(Dtype.int8), Skip(Dtype.uint8), Skip(Dtype.bool)],
-                },
-            ],
-        ),
-    ),
-
     'split': dict(
         name=["split"],
         tensor_para=dict(

diff --git a/DIOPI-IMPL/camb/functions/tril.cpp b/DIOPI-IMPL/camb/functions/tril.cpp
@@ -0,0 +1,39 @@
+/**
+ * @file
+ * @author DeepLink
+ * @copyright  (c) 2023, DeepLink.
+ */
+
+#include <diopi/functions.h>
+#include "../cnnl_helper.hpp"
+#include "../common/common.hpp"
+
+namespace impl {
+namespace camb {
+
+extern "C" {
+
+DIOPI_API diopiError_t diopiTril(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, int64_t diagonal) {
+    cnnlHandle_t handle = cnnlHandlePool.get(ctx);
+    DiopiTensor input_tensor(input);
+    DiopiTensor out_tensor(out);
+
+    std::vector<DiopiTensor*> pTensors{&input_tensor};
+    DIOPI_CALL(autoCastTensorType(ctx, pTensors, {diopi_dtype_int8, diopi_dtype_int16, diopi_dtype_int32, diopi_dtype_float16, diopi_dtype_float32}));
+    DiopiTensor input_tensor_tmp = *pTensors[0];
+    DiopiTensor out_tensor_tmp = out_tensor;
+    DIOPI_CALL(dataTypeCast(ctx, out_tensor_tmp, input_tensor_tmp.dtype()));
+
+    CnnlTensorDesc input_desc(input_tensor_tmp, CNNL_LAYOUT_ARRAY);
+    CnnlTensorDesc out_desc(out_tensor_tmp, CNNL_LAYOUT_ARRAY);
+
+    DIOPI_CALLCNNL(cnnlTri(handle, static_cast<int>(diagonal), false, input_desc.get(), input_tensor_tmp.data(), out_desc.get(), out_tensor_tmp.data()));
+    DIOPI_CALL(dataTypeCast(ctx, out_tensor, out_tensor_tmp));
+
+    return diopiSuccess;
+}
+
+}  // extern "C"
+
+}  // namespace camb
+}  // namespace impl