[onert] Support constant input BatchMatMul (#15033)

This commit adds support for constant input BatchMatMul.
It includes BatchMatMul constant input tests.

ONE-DCO-1.0-Signed-off-by: Hyeongseok Oh <hseok82.oh@samsung.com>
Co-authored-by: SeungHui Youn <61981457+zetwhite@users.noreply.github.com>

Author: hseok-oh

runtime/compute/cker/include/cker/operation/BatchMatMul.h

@@ -44,7 +44,8 @@ class BatchMatMul
   /**
    * @brief   Prepare temporary area for calculation
    */
-  void prepare(const Shape &lhs_shape, const Shape &rhs_shape, bool adj_x, bool adj_y)
+  void prepare(const Shape &lhs_shape, const Shape &rhs_shape, bool adj_x, bool adj_y,
+               bool rhs_const)
   {
     if (adj_x)
     {
@@ -75,18 +76,19 @@ class BatchMatMul
 
       _temp_rhs.resize(_temp_rhs_shape.FlatSize());
     }
+
+    _rhs_constant = rhs_const;
   }
 
   void operator()(const Shape &lhs_shape, const float *lhs_data, const Shape &rhs_shape,
                   const float *rhs_data, bool adj_x, bool adj_y, const Shape & /*output_shape*/,
                   float *output_data)
   {
-    // Assume lhs and rhs is not constant
-    // TODO Handle constant input
-
-    if (!adj_y)
+    // Don't need transpose if rhs is constant and already transposed
+    if (!adj_y && !(_rhs_constant && _rhs_transposed))
     {
       transposeRowsCols(rhs_shape, rhs_data, _temp_rhs_shape, _temp_rhs.data());
+      _rhs_transposed = true;
     }
 
     if (adj_x)
@@ -144,6 +146,8 @@ class BatchMatMul
   Shape _temp_lhs_shape;
   std::vector<float> _temp_rhs;
   Shape _temp_rhs_shape;
+  bool _rhs_constant = false;
+  bool _rhs_transposed = false;
 };
 
 } // namespace cker

runtime/compute/cker/include/cker/operation/Einsum.h

@@ -903,7 +903,8 @@ class Einsum
 
     // LaunchBatchMatMul::Launch(lhs, rhs, adj_x, adj_y, bcast, &output_reshaped);
     BatchMatMul batchMatMul;
-    batchMatMul.prepare(lhs.shape, rhs.shape, adj_x, adj_y);
+    // Set rhs is not constant: don't use optimization
+    batchMatMul.prepare(lhs.shape, rhs.shape, adj_x, adj_y, false);
     batchMatMul(lhs.shape, lhs.base<float>(), rhs.shape, rhs.base<float>(), adj_x, adj_y,
                 output_reshaped.shape, output_reshaped.base<float>());
   }

runtime/onert/backend/cpu/ops/BatchMatMulLayer.cc

@@ -39,7 +39,7 @@ void BatchMatMulLayer::batchMatMulFloat32()
 
   // TODO implement for constant input
 
-  batchmatmul_kernel.prepare(lhs_shape, rhs_shape, _adj_x, _adj_y);
+  batchmatmul_kernel.prepare(lhs_shape, rhs_shape, _adj_x, _adj_y, _rhs->is_constant());
   batchmatmul_kernel(lhs_shape, getBuffer<float>(_lhs), rhs_shape, getBuffer<float>(_rhs), _adj_x,
                      _adj_y, output_shape, getBuffer<float>(_output));
 }

runtime/onert/core/src/ir/OperationValidator.cc

@@ -117,8 +117,10 @@ void OperationValidator::visit(const operation::BatchMatMul &node)
   const auto rhs_index(node.getInputs().at(operation::BatchMatMul::Input::RHS));
   const auto output_index(node.getOutputs().at(0));
 
-  // Constant lhs and rhs is not implemented yet
-  OP_REQUIRES(!isConstant(lhs_index) && !isConstant(rhs_index));
+  // RHS can be constant, but LHS is not constant
+  // If one of inputs is constant, it must be RHS
+  // If two inputs are constant, BatchMatMul is optimized into constant by compiler
+  OP_REQUIRES(!isConstant(lhs_index));
 
   // Allow hybrid quantization (lhs: float / rhs: qint8 / out: float)
   OP_REQUIRES(isValidType(

runtime/tests/nnfw_api/src/GenModelTests/one_op_tests/BatchMatMul.test.cc

@@ -48,3 +48,40 @@ TEST_F(GenModelTest, neg_OneOp_BatchMatMul_InvalidType)
 
   SUCCEED();
 }
+
+TEST_F(GenModelTest, OneOp_BatchMatMul_Const)
+{
+  CircleGen cgen;
+  int lhs = cgen.addTensor({{1, 2, 3}, circle::TensorType::TensorType_FLOAT32});
+  std::vector<float> const_data{7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18};
+  uint32_t rhs_buf = cgen.addBuffer(const_data);
+  int rhs = cgen.addTensor({{1, 3, 4}, circle::TensorType::TensorType_FLOAT32, rhs_buf});
+  int out = cgen.addTensor({{1, 2, 4}, circle::TensorType::TensorType_FLOAT32});
+  cgen.addOperatorBatchMatMul({{lhs, rhs}, {out}}, false, false);
+  cgen.setInputsAndOutputs({lhs}, {out});
+  _context = std::make_unique<GenModelTestContext>(cgen.finish());
+  _context->addTestCase(TestCaseData{}
+                          .addInput<float>({1, 2, 3, 4, 5, 6})
+                          .addOutput<float>({74, 80, 86, 92, 173, 188, 203, 218}));
+  _context->setBackends({"cpu"});
+
+  SUCCEED();
+}
+
+TEST_F(GenModelTest, neg_OneOp_BatchMatMul_InvalidConst)
+{
+  // LHS constant is not allowed
+  CircleGen cgen;
+  std::vector<float> const_data{1, 2, 3, 4, 5, 6};
+  uint32_t lhs_buf = cgen.addBuffer(const_data);
+  int lhs = cgen.addTensor({{1, 2, 3}, circle::TensorType::TensorType_FLOAT32, lhs_buf});
+  int rhs = cgen.addTensor({{1, 3, 4}, circle::TensorType::TensorType_FLOAT32});
+  int out = cgen.addTensor({{1, 2, 4}, circle::TensorType::TensorType_FLOAT32});
+  cgen.addOperatorBatchMatMul({{lhs, rhs}, {out}}, false, false);
+  cgen.setInputsAndOutputs({lhs}, {out});
+  _context = std::make_unique<GenModelTestContext>(cgen.finish());
+  _context->setBackends({"cpu"});
+  _context->expectFailModelLoad();
+
+  SUCCEED();
+}