[mir] New interface for PadOp (#6708)

* Remove redundant `num_dims` parameter of constructor.
* Make padding arguments consistent with other operations (two vectors instead of vector of pairs).
* Change the type of the value to be used for padding from `Scalar` to `float` (it is `float` in all supported input frameworks).

Signed-off-by: Sergei Barannikov <s.barannikov@samsung.com>

diff --git a/compiler/mir/include/mir/ops/PadOp.h b/compiler/mir/include/mir/ops/PadOp.h
index 36d3fa2..fb4598b 100644 (file)
--- a/compiler/mir/include/mir/ops/PadOp.h
+++ b/compiler/mir/include/mir/ops/PadOp.h
@@ -25,13 +25,11 @@ namespace mir
 namespace ops
 {
 
-/**
- * @brief Pad operation class
- */
 class PadOp : public Operation
 {
 public:
   /**
+   * TODO Remove this constructor after eliminating all its uses.
    * @brief Class for Pad operation in modelIR
    * @param arg The input
    * @param num_dims Number of dimensions
@@ -40,39 +38,76 @@ public:
    */
   PadOp(Output *arg, int32_t num_dims, std::vector<std::pair<int32_t, int32_t>> paddings,
         const Scalar &scalar_value)
-      : Operation(Type::pad, {arg}), _num_dims(num_dims), _paddings(std::move(paddings)),
-        _scalar_value(scalar_value)
+      : Operation(Type::pad, {arg})
+
   {
+    _padding_before.resize(num_dims);
+    _padding_after.resize(num_dims);
+    for (int i = 0; i < num_dims; ++i)
+    {
+      _padding_before[i] = paddings[i].first;
+      _padding_after[i] = paddings[i].second;
+    }
+    _padding_value = scalar_value.get<float>();
+    inferOutputShapes();
+  }
+
+  /// @param arg The tensor to be padded.
+  /// @param padding_before The padding to be added before the tensor.
+  /// @param padding_after The padding to be added after the tensor.
+  /// @param padding_value The value to be used for padding.
+  PadOp(Output *arg, std::vector<std::int32_t> padding_before,
+        std::vector<std::int32_t> padding_after, float padding_value)
+      : Operation(Type::pad, {arg}), _padding_before(std::move(padding_before)),
+        _padding_after(std::move(padding_after)), _padding_value(padding_value)
+  {
+    assert(_padding_before.size() == _padding_after.size());
     inferOutputShapes();
   }
 
   Operation *copyWithInputs(const std::vector<Output *> &inputs) override
   {
-    return new PadOp(inputs[0], _num_dims, _paddings, _scalar_value);
+    return new PadOp(inputs[0], _padding_before, _padding_after, _padding_value);
   }
 
   /**
+   * TODO Remove this method after eliminating all its uses.
    * @param dim Dimension number
    * @return Pair of paddings for dimension
    */
-  std::pair<int32_t, int32_t> getPaddingForDim(int dim) const { return _paddings[dim]; }
+  std::pair<int32_t, int32_t> getPaddingForDim(int dim) const
+  {
+    return {_padding_before[dim], _padding_after[dim]};
+  }
 
   /**
+   * TODO Remove this method after eliminating all its uses.
    * @return Number of dimensions
    */
-  int getNumDim() const { return _num_dims; }
+  int getNumDim() const { return getInputShape(0).rank(); }
 
   /**
+   * TODO Remove this method after eliminating all its uses.
    * @return Scalar value
    */
-  Scalar getScalar() const { return _scalar_value; }
+  Scalar getScalar() const
+  {
+    return Scalar(reinterpret_cast<const char *>(&_padding_value), DataType::FLOAT32,
+                  sizeof(float));
+  }
+
+  const std::vector<std::int32_t> &getPaddingBefore() const { return _padding_before; }
+
+  const std::vector<std::int32_t> &getPaddingAfter() const { return _padding_after; }
+
+  float getPaddingValue() const { return _padding_value; }
 
 private:
   void inferOutputShapes();
 
-  std::vector<std::pair<int32_t, int32_t>> _paddings;
-  int32_t _num_dims;
-  Scalar _scalar_value;
+  std::vector<std::int32_t> _padding_before;
+  std::vector<std::int32_t> _padding_after;
+  float _padding_value;
 };
 
 } // namespace ops

diff --git a/compiler/mir/src/ops/PadOp.cpp b/compiler/mir/src/ops/PadOp.cpp
index f82bc00..75d8662 100644 (file)
--- a/compiler/mir/src/ops/PadOp.cpp
+++ b/compiler/mir/src/ops/PadOp.cpp
@@ -26,13 +26,10 @@ void PadOp::inferOutputShapes()
   const Shape &input_shape = getInputShape(0);
   int32_t num_dims = input_shape.rank();
 
-  assert(num_dims == getNumDim());
-
   Shape out_shape(num_dims);
   for (int32_t dim = 0; dim < num_dims; ++dim)
   {
-    std::pair<int32_t, int32_t> padding = getPaddingForDim(dim);
-    out_shape.dim(dim) = padding.first + input_shape.dim(dim) + padding.second;
+    out_shape.dim(dim) = _padding_before[dim] + input_shape.dim(dim) + _padding_after[dim];
   }
 
   setOutputShape(0, out_shape);