add a16 matmul multi_core

Change-Id: I10a9097ee52e324555f4a505ce18d7fe9b665803
sophgo · Nov 30, 2024 · 0fb281c · 0fb281c
1 parent c44a4a4
commit 0fb281c
Show file tree

Hide file tree

Showing 22 changed files with 27 additions and 5 deletions.
diff --git a/include/tpu_mlir/Backend/BM168x/Param.h b/include/tpu_mlir/Backend/BM168x/Param.h
@@ -1266,6 +1266,7 @@ typedef struct a16_matmul_spec {
   int weight_bits;
   bool has_zp;
   int q_group_size;
+  int use_multi_core;
 } a16_matmul_spec_t;
 
 typedef struct batch_matmul_common_spec {

diff --git a/include/tpu_mlir/Support/Module.h b/include/tpu_mlir/Support/Module.h
@@ -191,6 +191,7 @@ bool IsHdimIsBatch(Operation *Op);
 bool IsHdimIsBatch(Value value);
 bool isOpInCoreMatch(Operation *Op);
 bool isOpInCoreParallel(Operation *Op);
+bool isOpInGroupParallel(Operation *Op);
 bool isOpInDevParallel(Operation *Op);
 bool isOpInBlock(Operation *op);
 FuncOp getFuncOp(ModuleOp module, StringRef func_name);

diff --git a/lib/Dialect/Tpu/Interfaces/BM1684X/A16MatMul.cpp b/lib/Dialect/Tpu/Interfaces/BM1684X/A16MatMul.cpp
@@ -26,8 +26,14 @@ void tpu::A16MatMulOp::codegen_global_bm1684x() {
   auto op = getOperation();
   auto input_spec = BM168x::get_input_spec(op);
   auto output_spec = BM168x::get_output_spec(op);
-  BM168x::call_global_func("backend_api_a16_matmul_global", &spec, sizeof(spec),
-                           input_spec->data(), output_spec->data());
+  if (supportMultiCore(*this)){
+    spec.use_multi_core = 1;
+  } else {
+    spec.use_multi_core = 0;
+  }
+  BM168x::call_global_func("backend_api_a16_matmul_global", &spec,
+                           sizeof(spec), input_spec->data(),
+                           output_spec->data());
 }
 
 // ======================================

diff --git a/lib/Dialect/Tpu/Interfaces/Common/A16MatMul.cpp b/lib/Dialect/Tpu/Interfaces/Common/A16MatMul.cpp
@@ -493,4 +493,4 @@ ArrayAttr tpu::A16MatMulOp::getIndexingMaps() {
   return Builder(ctx).getAffineMapArrayAttr(indexingMaps);
 }
 
-bool tpu::A16MatMulOp::support_multi_core() { return module::isSG2380(); }
+bool tpu::A16MatMulOp::support_multi_core() { return (module::isSG2380() || module::isBM1690Family()) && !module::isOpInGroupParallel(*this);}
diff --git a/lib/Support/Module.cpp b/lib/Support/Module.cpp
@@ -826,6 +826,17 @@ bool isOpInCoreParallel(Operation *Op) {
   return false;
 }
 
+bool isOpInGroupParallel(Operation *Op) {
+  if (Op == nullptr) {
+    return false;
+  }
+  auto parent = Op->getParentOp();
+  if (isa_and_nonnull<tpu::GroupParallelOp>(parent)) {
+    return true;
+  }
+  return false;
+}
+
 // op in [CoreBegin, CoreEnd]
 bool isOpInCoreMatch(Operation *op) {
   while (!op->use_empty()) {

diff --git a/third_party/nntoolchain/README.md b/third_party/nntoolchain/README.md
@@ -17,6 +17,7 @@ cp bmcompiler/libbackend/libbackend_1684.so /workspace/tpu-mlir/third_party/nnto
 
 #bm1684x sha256:  b9653fd714815e33546701730698c4f14114b5dd
 # - yolov5s perf revive: 5fd5cfea836a56e91faf59ae669b1865959019a4
+# + a16matmul multi_core: ead92658fce9ad437a71a3a873600d968ded14db
 cd TPU1686
 source  scripts/envsetup.sh bm1684x
 debug: rebuild_backend_lib_cmodel
@@ -29,6 +30,7 @@ cp build/firmware_core/libfirmware_core.a /workspace/tpu-mlir/third_party/nntool
 
 #bm1688 sha256: 8453b87d6b1ad8858692f8ca4cd2441e7b677f1e
 # - MDSR_x3/x4 perf revive: 253d8e58e7b2764d625bfcdd9d049d0184a0f602
+# + a16matmul multi_core: ead92658fce9ad437a71a3a873600d968ded14db
 cd TPU1686
 source  scripts/envsetup.sh bm1686
 debug: rebuild_backend_lib_cmodel
@@ -42,6 +44,7 @@ cp build/firmware_core/libfirmware_core.a /workspace/tpu-mlir/third_party/nntool
 #bm1690 sha256: 90428e6d47c53c73fab8f846a1f579af456862aa
 # + MaxPoolingWithMask optimize: 85e07f6a2e6059f84b1b591b2df40fe7e88f21dc
 # + RMSNorm global layer bugfix: d100aad8b3b7d6b9d262b5aaee7e1e0f46daef43 (not merged)
+# + a16matmul multi_core: ead92658fce9ad437a71a3a873600d968ded14db
 cd TPU1686
 source  scripts/envsetup.sh sg2260
 debug: rebuild_backend_lib_cmodel
@@ -53,15 +56,15 @@ unset EXTRA_CONFIG && rebuild_firmware
 cp build/firmware_core/libfirmware_core.so /workspace/tpu-mlir/third_party/nntoolchain/lib/libbm1690_kernel_module.so
 cp build/firmware_core/libfirmware_core.a /workspace/tpu-mlir/third_party/nntoolchain/lib/libbm1690_kernel_module.a
 
-#sg2380 sha256: 8453b87d6b1ad8858692f8ca4cd2441e7b677f1e
+#sg2380 sha256: ead92658fce9ad437a71a3a873600d968ded14db
 cd TPU1686
 source  scripts/envsetup.sh sg2380
 debug: rebuild_backend_lib_cmodel
 elease: unset EXTRA_CONFIG && rebuild_backend_lib_release_cmodel
 cp build/backend_api/libbackend_sg2380.so /workspace/tpu-mlir/third_party/nntoolchain/lib/libbackend_sg2380.so
 cp build_runtime/firmware_core/libcmodel_firmware.so /workspace/tpu-mlir/third_party/nntoolchain/lib/libcmodel_sg2380.so
 
-#mars3 sha256:  68a4781b8f0ab6bb38eab38e4d1fd9c6f8a426b4
+#mars3 sha256:  ead92658fce9ad437a71a3a873600d968ded14db
 cd TPU1686
 source  scripts/envsetup.sh mars3
 debug: rebuild_backend_lib_cmodel