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/TosaDeserialize.cpp | 22 ++++++++++++++++++++++
 src/TosaSerialize.cpp   | 21 +++++++++++++++++++++
 2 files changed, 43 insertions(+)

(limited to 'src')

diff --git a/src/TosaDeserialize.cpp b/src/TosaDeserialize.cpp
index 70a19a7..9e646c7 100644
--- a/src/TosaDeserialize.cpp
+++ b/src/TosaDeserialize.cpp
@@ -87,6 +87,9 @@ mlir::LogicalResult BuildTensorType(mlir::OpBuilder *op_builder,
   case DType_UINT16:
     element_type = op_builder->getIntegerType(16, false);
     break;
+  case DType_SHAPE:
+    element_type = op_builder->getIntegerType(64);
+    break;
   default:
     llvm::errs() << "ERROR: unknown type " << EnumNamesDType()[ts->GetDtype()]
                  << "\n";
@@ -859,6 +862,25 @@ TosaMlirOperatorBuilder::build<Op_PAD>(TosaSerializationOperator *op) const {
   return std::vector<mlir::Value>({mlir_op->getResult(0)});
 }
 
+template <>
+std::vector<mlir::Value>
+TosaMlirOperatorBuilder::build<Op_DIM>(TosaSerializationOperator *op) const {
+  mlir::Value input_val = tensor_map->at(op->GetInputTensorNames()[0]);
+  mlir::RankedTensorType output_type =
+      tensor_type_map->at(op->GetOutputTensorNames()[0]);
+
+  assert(op->GetAttributeType() ==
+         Attribute_AxisAttribute); // double check attribute type
+  TosaAxisAttribute *attr =
+      static_cast<TosaAxisAttribute *>(op->GetAttribute());
+  auto axis = op_builder->getI32IntegerAttr(attr->axis());
+
+  mlir::Operation *mlir_op =
+      op_builder->create<mlir::tosa::DimOp>(loc, output_type, input_val, axis);
+  block->push_back(mlir_op);
+  return std::vector<mlir::Value>({mlir_op->getResult(0)});
+}
+
 template <>
 std::vector<mlir::Value> TosaMlirOperatorBuilder::build<Op_TRANSPOSE>(
     TosaSerializationOperator *op) const {
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