From ea49f62f7ab81750f19bef011683164fe9bd4080 Mon Sep 17 00:00:00 2001
From: Won Jeon <won.jeon@arm.com>
Date: Thu, 10 Aug 2023 23:00:22 -0700
Subject: Add DIM operator and its serialization/deserialization to TOSA MLIR
 translator

Signed-off-by: Won Jeon <won.jeon@arm.com>
Change-Id: Idf45f0a68f551039cc40a69d90aa5c53098bc238
---
 src/TosaSerialize.cpp | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

(limited to 'src/TosaSerialize.cpp')

diff --git a/src/TosaSerialize.cpp b/src/TosaSerialize.cpp
index 17a693e..741597b 100644
--- a/src/TosaSerialize.cpp
+++ b/src/TosaSerialize.cpp
@@ -86,6 +86,9 @@ static DType Type2DType(mlir::Type element_type) {
     return DType_INT32;
   } else if (element_type.isInteger(48)) {
     return DType_INT48;
+  } else if (element_type.isInteger(64)) {
+    // shape treated as integer with bitwidth 64 for now
+    return DType_SHAPE;
   }
   // boolean in MLIR treated as integer with bitwidth 1
   else if (element_type.isInteger(1)) {
@@ -974,6 +977,24 @@ TosaSerializationOperatorBuilder::build<mlir::tosa::PadOp>(
   return tyop;
 }
 
+template <>
+TosaSerializationOperator *
+TosaSerializationOperatorBuilder::build<mlir::tosa::DimOp>(
+    mlir::Operation &op) const {
+  std::string input_name = GetTensorName(op.getOperand(0));
+  std::string output_name = GetTensorName(op.getResult(0));
+
+  int32_t axis = op.getAttr("axis").dyn_cast<mlir::IntegerAttr>().getInt();
+  TosaAxisAttribute attribute(axis);
+
+  TosaSerializationOperator *tyop =
+      new TosaSerializationOperator(Op_DIM, Attribute_AxisAttribute, &attribute,
+                                    std::vector<std::string>{input_name},
+                                    std::vector<std::string>{output_name});
+
+  return tyop;
+}
+
 template <>
 TosaSerializationOperator *
 TosaSerializationOperatorBuilder::build<mlir::tosa::TransposeOp>(
-- 
cgit v1.2.1