Fix Div

Change-Id: I66ca38746e19299c544b27ad4feab004adfae5d3
sophgo · Feb 15, 2023 · 123549b · 123549b
1 parent 5315aa6
commit 123549b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/lib/Dialect/Top/Interfaces/Div.cpp b/lib/Dialect/Top/Interfaces/Div.cpp
@@ -16,9 +16,14 @@ int64_t top::DivOp::getFLOPs() { return module::getNumElements(getOutput()); }
 
 LogicalResult top::DivOp::init(InferenceParameter &p) {
   auto binary = new Binary();
+  auto lhs_shape =  module::getShape(getInputs()[0]);
+  auto rhs_shape = module::getShape(getInputs()[1]);
+  auto max_ndim = std::max(lhs_shape.size(), rhs_shape.size());
+  auto input0_shape = shape_expand_dim(lhs_shape, max_ndim);
+  auto input1_shape = shape_expand_dim(rhs_shape, max_ndim);
   (*binary)
-      .lhs(p.inputs[0], module::getShape(getInputs()[0]))
-      .rhs(p.inputs[1], module::getShape(getInputs()[1]))
+      .lhs(p.inputs[0], input0_shape)
+      .rhs(p.inputs[1], input1_shape)
       .dst(p.outputs[0], module::getShape(getOutput()))
       .do_relu(getDoRelu())
       .relu_limit(getReluLimit().convertToDouble())

diff --git a/lib/Dialect/Tpu/Interfaces/Common/Div.cpp b/lib/Dialect/Tpu/Interfaces/Common/Div.cpp
@@ -15,9 +15,14 @@
 
 LogicalResult tpu::DivOp::init(InferenceParameter &p) {
   auto binary = new Binary();
+  auto lhs_shape =  module::getShape(getInputs()[0]);
+  auto rhs_shape = module::getShape(getInputs()[1]);
+  auto max_ndim = std::max(lhs_shape.size(), rhs_shape.size());
+  auto input0_shape = shape_expand_dim(lhs_shape, max_ndim);
+  auto input1_shape = shape_expand_dim(rhs_shape, max_ndim);
   (*binary)
-      .lhs(p.inputs[0], module::getShape(getInputs()[0]))
-      .rhs(p.inputs[1], module::getShape(getInputs()[1]))
+      .lhs(p.inputs[0], input0_shape)
+      .rhs(p.inputs[1], input1_shape)
       .dst(p.outputs[0], module::getShape(getOutput()))
       .do_relu(getDoRelu())
       .relu_limit(getReluLimit().convertToDouble())