From 182ac7eb30b58993bb8d7682f6354cf0104d5a95 Mon Sep 17 00:00:00 2001
From: jiangchengcheng-on <jiangchengcheng3@h-partners.com>
Date: Fri, 9 Aug 2024 09:25:18 +0000
Subject: [PATCH] optimize mm api

Signed-off-by: jiangchengcheng-on <jiangchengcheng3@h-partners.com>
---
 impl/matmul/kernel_kfc.h | 38 ++++++++++++++++++++++----------------
 1 file changed, 22 insertions(+), 16 deletions(-)

diff --git a/impl/matmul/kernel_kfc.h b/impl/matmul/kernel_kfc.h
index bcf0a472..14e3f017 100644
--- a/impl/matmul/kernel_kfc.h
+++ b/impl/matmul/kernel_kfc.h
@@ -236,41 +236,46 @@ constexpr bool IsSharedMatmul()
     return !MM_CFG.enableInit;
 }
 template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE = C_TYPE,
-    const MatmulConfig& MM_CFG = CFG_NORM, class MM_CB = matmul::MatmulCallBackFunc<nullptr, nullptr, nullptr>>
+    const MatmulConfig& MM_CFG = CFG_NORM, class MM_CB = matmul::MatmulCallBackFunc<nullptr, nullptr, nullptr>,
+    MATMUL_POLICY_DEFAULT_OF(matmul::MatmulPolicy)>
 struct MatmulInstBase {
     __aicore__ inline MatmulInstBase(){};
 };
-template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE, const MatmulConfig& MM_CFG, class MM_CB>
-struct MatmulInstShared : MatmulInstBase<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB> {
+template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE, const MatmulConfig& MM_CFG, class MM_CB,
+    MATMUL_POLICY_TEMPLATE_OF(MATMUL_POLICY)>
+struct MatmulInstShared : MatmulInstBase<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB, MATMUL_POLICY> {
     __aicore__ inline MatmulInstShared(){};
-    matmul::MatmulService<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB> mm[1];
+    matmul::MatmulService<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB, MATMUL_POLICY> mm[1];
 };
-template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE, const MatmulConfig& MM_CFG, class MM_CB>
-struct MatmulInst : MatmulInstBase<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB> {
+template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE, const MatmulConfig& MM_CFG, class MM_CB,
+    MATMUL_POLICY_TEMPLATE_OF(MATMUL_POLICY)>
+struct MatmulInst : MatmulInstBase<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB, MATMUL_POLICY> {
     __aicore__ inline MatmulInst(){};
-    matmul::MatmulService<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB> mm[MIX_NUM];
+    matmul::MatmulService<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB, MATMUL_POLICY> mm[MIX_NUM];
 };
 
 template <bool SHARED, class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE, const MatmulConfig& MM_CFG,
-    class MM_CB>
+    class MM_CB, MATMUL_POLICY_TEMPLATE_OF(MATMUL_POLICY)>
 struct MatmulInstAux {
     __aicore__ inline MatmulInstAux(){};
 };
 
-template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE, const MatmulConfig& MM_CFG, class MM_CB>
-struct MatmulInstAux<true, A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB> {
+template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE, const MatmulConfig& MM_CFG, class MM_CB,
+    MATMUL_POLICY_TEMPLATE_OF(MATMUL_POLICY)>
+struct MatmulInstAux<true, A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB, MATMUL_POLICY> {
     __aicore__ inline MatmulInstAux(){};
-    using MATMUL = MatmulInstShared<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB>;
+    using MATMUL = MatmulInstShared<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB, MATMUL_POLICY>;
 };
 
-template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE, const MatmulConfig& MM_CFG, class MM_CB>
-struct MatmulInstAux<false, A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB> {
+template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE, const MatmulConfig& MM_CFG, class MM_CB,
+    MATMUL_POLICY_TEMPLATE_OF(MATMUL_POLICY)>
+struct MatmulInstAux<false, A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB, MATMUL_POLICY> {
     __aicore__ inline MatmulInstAux(){};
-    using MATMUL = MatmulInst<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB>;
+    using MATMUL = MatmulInst<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB, MATMUL_POLICY>;
 };
 
 template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE = C_TYPE, const MatmulConfig& MM_CFG = CFG_NORM,
-    class MM_CB = matmul::MatmulCallBackFunc<nullptr, nullptr, nullptr>>
+class MM_CB = matmul::MatmulCallBackFunc<nullptr, nullptr, nullptr>, MATMUL_POLICY_DEFAULT_OF(matmul::MatmulPolicy)>
 class MatmulServiceAux {
     using SrcT = typename A_TYPE::T;
     using SrcAT = typename A_TYPE::T;
@@ -281,7 +286,8 @@ class MatmulServiceAux {
 
 public:
     __aicore__ inline MatmulServiceAux() {}
-    typename MatmulInstAux<IsSharedMatmul<MM_CFG>(), A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB>::MATMUL mm;
+    typename MatmulInstAux<IsSharedMatmul<MM_CFG>(), A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, MM_CFG, MM_CB,
+        MATMUL_POLICY>::MATMUL mm;
 
     // stub functions for MatmulImpl
     __aicore__ inline void Init(TCubeTiling* cubeTiling, TPipe* tpipe = nullptr){};
-- 
Gitee