[loco] Introduce Softmax node in loco (#6492)

This commit will introduce `Softmax` node in `loco`

Signed-off-by: Seok NamKoong <sk.namkoong@samsung.com>

diff --git a/compiler/loco/include/loco/IR/CanonicalNodes.lst b/compiler/loco/include/loco/IR/CanonicalNodes.lst
index 771dc35..49808d3 100644 (file)
--- a/compiler/loco/include/loco/IR/CanonicalNodes.lst
+++ b/compiler/loco/include/loco/IR/CanonicalNodes.lst
@@ -32,3 +32,4 @@ CANONICAL_NODE(ReLU, ReLU)
 CANONICAL_NODE(ReLU6, ReLU6)
 CANONICAL_NODE(TensorConcat, TensorConcat)
 CANONICAL_NODE(TensorBiasAdd, BiasAdd<Domain::Tensor>)
+CANONICAL_NODE(TensorSoftmax, Softmax<Domain::Tensor>)

diff --git a/compiler/loco/include/loco/IR/Nodes.h b/compiler/loco/include/loco/IR/Nodes.h
index a858ad0..535136d 100644 (file)
--- a/compiler/loco/include/loco/IR/Nodes.h
+++ b/compiler/loco/include/loco/IR/Nodes.h
@@ -544,6 +544,34 @@ private:
 };
 
 /**
+ * @brief Computes softmax activations
+ */
+template <Domain D> class Softmax;
+
+/**
+* @brief Computes softmax activations for Tensor domain
+*/
+template <>
+class Softmax<Domain::Tensor> final
+    : public CanonicalNodeDef<CanonicalOpcode::TensorSoftmax, FixedArity<1>::Mixin>
+{
+public:
+  Softmax() = default;
+
+public:
+  Node *input(void) const { return at(0)->node(); }
+  void input(Node *node) { return at(0)->node(node); }
+
+  uint32_t axis(void) const { return _axis; }
+  void axis(uint32_t axis) { _axis = axis; }
+
+private:
+  uint32_t _axis = 0;
+};
+
+using TensorSoftmax = Softmax<Domain::Tensor>;
+
+/**
  * @brief Create a "Bias" from a "Tensor"
  *
  * BiasEncode currently requires a rank-1 tensor as its input.