ScaleOp support split to multi core

Change-Id: I574953ca809fac0ec8535fddb7f00cc12dfdba41
sophgo · Nov 21, 2024 · 4dff27c · 4dff27c
1 parent cd758f8
commit 4dff27c
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 0 deletions.
diff --git a/include/tpu_mlir/Dialect/Tpu/IR/TpuOps.td b/include/tpu_mlir/Dialect/Tpu/IR/TpuOps.td
@@ -2471,6 +2471,7 @@ def Tpu_Deconv3DOp: Tpu_Op<"Deconv3D",[
 def Tpu_ScaleOp: Tpu_Op<"Scale", [
   SupportFuseRelu, InOutSameShape,
   DeclareOpInterfaceMethods<LocalGenInterface>,
+  DeclareOpInterfaceMethods<IndexingMapsInterface>,
   DeclareOpInterfaceMethods<DynLocalGenInterface>]> {
   let summary = "Scale operator";
 

diff --git a/include/tpu_mlir/Interfaces/IndexingMapsInterface.h b/include/tpu_mlir/Interfaces/IndexingMapsInterface.h
@@ -16,5 +16,7 @@
 namespace tpu_mlir {
 
 mlir::ArrayAttr getBinaryIndexingMaps(mlir::Operation *op);
+mlir::AffineMap getBinaryMap(const mlir::AffineMap &refer_map,
+                             const llvm::ArrayRef<int64_t> &shape);
 
 } // namespace tpu_mlir
diff --git a/lib/Dialect/Tpu/Interfaces/Common/Scale.cpp b/lib/Dialect/Tpu/Interfaces/Common/Scale.cpp
@@ -93,4 +93,24 @@ LogicalResult tpu::ScaleOp::inference(InferenceParameter &p) {
   return success();
 }
 
+ArrayAttr tpu::ScaleOp::getIndexingMaps() {
+  MLIRContext *ctx = getContext();
+  auto out_shape = module::getShape(getOutput());
+  auto num_dims = out_shape.size();
+  SmallVector<AffineMap, 8> indexingMaps;
+  AffineMap outMap = AffineMap::getMultiDimIdentityMap(num_dims, ctx);
+  auto empty_map = AffineMap::get(num_dims, 0, ctx);
+  for (auto opd : getOperands()) {
+    if (module::isNone(opd)) {
+      indexingMaps.emplace_back(empty_map);
+      continue;
+    }
+    auto shape = module::getShape(opd);
+    auto map = getBinaryMap(outMap, shape);
+    indexingMaps.emplace_back(map);
+  }
+  indexingMaps.emplace_back(outMap);
+  return Builder(ctx).getAffineMapArrayAttr(indexingMaps);
+}
+
 bool tpu::ScaleOp::support_multi_core() { return false; }
diff --git a/lib/Interfaces/IndexingMapsInterface.cpp b/lib/Interfaces/IndexingMapsInterface.cpp
@@ -48,4 +48,23 @@ mlir::ArrayAttr getBinaryIndexingMaps(mlir::Operation *op) {
   return Builder(ctx).getAffineMapArrayAttr(indexingMaps);
 };
 
+mlir::AffineMap getBinaryMap(const mlir::AffineMap &refer_map,
+                             const llvm::ArrayRef<int64_t> &shape) {
+  if (refer_map.getNumDims() != shape.size()) {
+    llvm_unreachable("map dim is not the same");
+  }
+  auto ctx = refer_map.getContext();
+  auto one = mlir::getAffineConstantExpr(1, ctx);
+  auto num_dim = shape.size();
+  SmallVector<AffineExpr> index_v;
+  for (int i = 0; i < num_dim; i++) {
+    if (shape[i] == 1) {
+      index_v.push_back(one);
+    } else {
+      index_v.push_back(refer_map.getResults()[i]);
+    }
+  }
+  return AffineMap::get(num_dim, 0, index_v, ctx);
+}
+
 }; // namespace tpu_mlir