From 20ab82d00cbab2219cabeddf5d402e57e7403b28 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Mon, 5 Jul 2021 19:00:11 +0800
Subject: [PATCH] fix(tee): fix tee crash

GitOrigin-RevId: 379f970c8777f6ba1e6acce5ab46ff1dc6954a1c
---
 dnn/src/aarch64/matrix_mul/algos.cpp                        | 6 ++++++
 dnn/src/aarch64/matrix_mul/fp32/strategy.cpp                | 6 ++++++
 .../direct_kernels/f32_direct_nchw_nchw44_kern_common.h     | 4 ++++
 3 files changed, 16 insertions(+)

diff --git a/dnn/src/aarch64/matrix_mul/algos.cpp b/dnn/src/aarch64/matrix_mul/algos.cpp
index 0b640d328..13104c8cd 100644
--- a/dnn/src/aarch64/matrix_mul/algos.cpp
+++ b/dnn/src/aarch64/matrix_mul/algos.cpp
@@ -908,6 +908,9 @@ bool MatrixMulImpl::AlgoInt8x8x16MK4_16x12x4::preferred(
     return false;
 #else
     auto arch = cpuinfo_get_current_core()->uarch;
+#ifdef __IN_TEE_ENV__
+    arch = cpuinfo_uarch_unknown;
+#endif
     bool little_core = arch == cpuinfo_uarch_cortex_a53 ||
                        arch == cpuinfo_uarch_cortex_a55;
     return little_core;
@@ -986,6 +989,9 @@ bool MatrixMulImpl::AlgoInt8x8x16MK4_4x4x8::preferred(
     return false;
 #else
     auto arch = cpuinfo_get_current_core()->uarch;
+#ifdef __IN_TEE_ENV__
+    arch = cpuinfo_uarch_unknown;
+#endif
     bool little_core = arch == cpuinfo_uarch_cortex_a53 ||
                        arch == cpuinfo_uarch_cortex_a55;
     return !little_core;
diff --git a/dnn/src/aarch64/matrix_mul/fp32/strategy.cpp b/dnn/src/aarch64/matrix_mul/fp32/strategy.cpp
index 14c1238bd..9b5d51384 100644
--- a/dnn/src/aarch64/matrix_mul/fp32/strategy.cpp
+++ b/dnn/src/aarch64/matrix_mul/fp32/strategy.cpp
@@ -181,6 +181,9 @@ void sgemm_8x12::kern(const float* packA, const float* packB, size_t M,
                                            is_first_k);
 #else
     auto arch = cpuinfo_get_current_core()->uarch;
+#ifdef __IN_TEE_ENV__
+    arch = cpuinfo_uarch_unknown;
+#endif
     if (arch == cpuinfo_uarch_cortex_a53) {
         sgemm_8x12_helper<matmul_general_8x12_a53>(packA, packB, M, N, K, C,
                                                    LDC, is_first_k);
@@ -274,6 +277,9 @@ void sgemm_mk4_8x12::kern(const float* packA, const float* packB, size_t M,
                                            is_first_k);
 #else
     auto arch = cpuinfo_get_current_core()->uarch;
+#ifdef __IN_TEE_ENV__
+    arch = cpuinfo_uarch_unknown;
+#endif
     if (arch == cpuinfo_uarch_cortex_a53) {
         sgemm_mk4_8x12_helper<matmul_mk4_8x12_a53>(packA, packB, M, N, K, C,
                                                    LDC, is_first_k);
diff --git a/dnn/src/arm_common/conv_bias/fp32/direct_kernels/f32_direct_nchw_nchw44_kern_common.h b/dnn/src/arm_common/conv_bias/fp32/direct_kernels/f32_direct_nchw_nchw44_kern_common.h
index 9184f278c..8090c5e5c 100644
--- a/dnn/src/arm_common/conv_bias/fp32/direct_kernels/f32_direct_nchw_nchw44_kern_common.h
+++ b/dnn/src/arm_common/conv_bias/fp32/direct_kernels/f32_direct_nchw_nchw44_kern_common.h
@@ -851,6 +851,10 @@ struct ConvDirectFp32NchwNchw44<bias_mode, Op, 3, 2> {
 #else
         auto arch_tag = CpuTag::A7_TAG;
 #endif
+
+#ifdef __IN_TEE_ENV__
+        arch_tag = CpuTag::DEFAULT_CPU_TAG;
+#endif
         if (arch_tag == CpuTag::A7_TAG) {
             for (int oc_idx = 0; oc_idx < oc_end; oc_idx += big_oc_step) {
                 const int weight_offset = oc_idx * ic * fh * fw;
-- 
GitLab