[tosa_mlir_translator] Add Shape Ops Support

- Add serialization/deserialization of tosa shape ops.
- Changed TileOp's multiples from attribute to shape input.
- Change 'shape' attribute of RESHAPE to an input

Signed-off-by: Tai Ly <tai.ly@arm.com>
Change-Id: I4329e621fd7637b1a3491c195fbda77d2a0ad23a

1 files changed, 92 insertions, 16 deletions

diff --git a/src/TosaSerialize.cpp b/src/TosaSerialize.cpp
index 2d038f0..4f2c358 100644
--- a/src/TosaSerialize.cpp
+++ b/src/TosaSerialize.cpp
@@ -1,5 +1,5 @@
 
-// Copyright (c) 2020-2023, ARM Limited.
+// Copyright (c) 2020-2024, ARM Limited.
 //
 //    Licensed under the Apache License, Version 2.0 with LLVM Exceptions
 //    (the "License"); you may not use this file except in compliance with
@@ -93,9 +93,6 @@ 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)) {
@@ -608,6 +605,57 @@ BUILD_OP_ELEMENTWISE_UNARY(Tanh, TANH)
 BUILD_OP_ELEMENTWISE_UNARY(Identity, IDENTITY)
 BUILD_OP_ELEMENTWISE_UNARY(Cast, CAST)
 
+BUILD_OP_ELEMENTWISE_BINARY(AddShape, ADD_SHAPE)
+BUILD_OP_ELEMENTWISE_BINARY(SubShape, SUB_SHAPE)
+BUILD_OP_ELEMENTWISE_BINARY(MulShape, MUL_SHAPE)
+BUILD_OP_ELEMENTWISE_BINARY(DivShape, DIV_SHAPE)
+
+template <>
+TosaSerializationOperator *
+TosaSerializationOperatorBuilder::build<mlir::tosa::ConstShapeOp>(
+    mlir::Operation &op) const {
+  std::string output_name = GetTensorName(op.getResult(0));
+  TosaSerializationTensor *ts =
+      block_builder->GetBlock()->GetTensorByName(output_name);
+  if (!ts) {
+    op.emitOpError(
+        "ERROR: serialization tensor must be built before building operator");
+    return nullptr;
+  }
+
+  // Update tensor.data array with Const value attribute
+  mlir::Attribute value_attr = op.getAttr("value");
+  if (!value_attr) {
+    op.emitOpError("ERROR: tosa.const_shape doesn't have value");
+    return nullptr;
+  }
+  assert(ts->GetDtype() == DType::DType_SHAPE);
+  std::vector<uint8_t> u8_data;
+
+  std::vector<int64_t> data;
+  auto dense_attr = op.getAttr(llvm::StringRef("value"))
+                        .dyn_cast<mlir::DenseIntElementsAttr>();
+  if (!dense_attr) {
+    op.emitOpError("Unknown const attribute");
+    return nullptr;
+  }
+
+  for (auto valueIt : dense_attr.getValues<mlir::APInt>()) {
+    int64_t val = valueIt.getSExtValue();
+    data.push_back(val);
+  }
+
+  TosaSerializationHandler::ConvertI64toU8(data, u8_data);
+
+  ts->SetData(u8_data);
+
+  TosaSerializationOperator *tyop = new TosaSerializationOperator(
+      Op_CONST_SHAPE, Attribute_NONE, nullptr, std::vector<std::string>{},
+      std::vector<std::string>{output_name});
+
+  return tyop;
+}
+
 template <>
 TosaSerializationOperator *
 TosaSerializationOperatorBuilder::build<mlir::tosa::ConstOp>(
@@ -1097,6 +1145,25 @@ TosaSerializationOperatorBuilder::build<mlir::tosa::ConcatOp>(
 
 template <>
 TosaSerializationOperator *
+TosaSerializationOperatorBuilder::build<mlir::tosa::ConcatShapeOp>(
+    mlir::Operation &op) const {
+  std::vector<std::string> inputs;
+  for (uint32_t i = 0; i < op.getNumOperands(); i++) {
+    std::string input_name = GetTensorName(op.getOperand(i));
+    inputs.push_back(input_name);
+  }
+
+  std::string output_name = GetTensorName(op.getResult(0));
+
+  TosaSerializationOperator *tyop = new TosaSerializationOperator(
+      Op_CONCAT_SHAPE, Attribute_NONE, nullptr, inputs,
+      std::vector<std::string>{output_name});
+
+  return tyop;
+}
+
+template <>
+TosaSerializationOperator *
 TosaSerializationOperatorBuilder::build<mlir::tosa::NegateOp>(
     mlir::Operation &op) const {
   std::string input_name = GetTensorName(op.getOperand(0));
@@ -1122,14 +1189,12 @@ TosaSerializationOperator *
 TosaSerializationOperatorBuilder::build<mlir::tosa::ReshapeOp>(
     mlir::Operation &op) const {
   std::string input_name = GetTensorName(op.getOperand(0));
+  std::string shape_name = GetTensorName(op.getOperand(1));
   std::string output_name = GetTensorName(op.getResult(0));
-  auto shape = getDenseI64ArrayAttr<int>(op.getAttr("new_shape"));
-
-  TosaReshapeAttribute attribute(shape);
 
   TosaSerializationOperator *tyop = new TosaSerializationOperator(
-      Op_RESHAPE, Attribute_ReshapeAttribute, &attribute,
-      std::vector<std::string>{input_name},
+      Op_RESHAPE, Attribute_NONE, nullptr,
+      std::vector<std::string>{input_name, shape_name},
       std::vector<std::string>{output_name});
 
   return tyop;
@@ -1255,15 +1320,13 @@ template <>
 TosaSerializationOperator *
 TosaSerializationOperatorBuilder::build<mlir::tosa::TileOp>(
     mlir::Operation &op) const {
-  std::string input_name = GetTensorName(op.getOperand(0));
+  std::string input0_name = GetTensorName(op.getOperand(0));
+  std::string input1_name = GetTensorName(op.getOperand(1));
   std::string output_name = GetTensorName(op.getResult(0));
-  auto multiples = getDenseI64ArrayAttr<int>(op.getAttr("multiples"));
-
-  TosaTileAttribute attribute(multiples);
 
   TosaSerializationOperator *tyop = new TosaSerializationOperator(
-      Op_TILE, Attribute_TileAttribute, &attribute,
-      std::vector<std::string>{input_name},
+      Op_TILE, Attribute_NONE, nullptr,
+      std::vector<std::string>{input0_name, input1_name},
       std::vector<std::string>{output_name});
 
   return tyop;
@@ -1996,6 +2059,20 @@ TosaSerializationBlockBuilder::BuildTosaSerializationTensor(
     return nullptr;
   }
 
+  // handling of tosa.shape values
+  if (auto shape_ty = val.getType().dyn_cast<mlir::tosa::shapeType>()) {
+    auto rank = shape_ty.getRank();
+    std::vector<int32_t> shape;
+    if (rank > 0) {
+      shape.push_back(rank);
+    }
+    ts = new TosaSerializationTensor(name,
+                                     /* shape = */ shape,
+                                     /* type = */ DType::DType_SHAPE,
+                                     /* data = */ std::vector<uint8_t>());
+    return ts;
+  }
+
   auto ttype = val.getType().dyn_cast<mlir::TensorType>();
   if (!ttype) {
     llvm::errs() << "TOSA serialization, supplied value is not of TensorType\n";
@@ -2012,7 +2089,6 @@ TosaSerializationBlockBuilder::BuildTosaSerializationTensor(
         shape.push_back(0); // size of 0 represents dynamic dimension
       } else {
         auto dim = shaped.getDimSize(idx);
-        assert(dim > 0);
         shape.push_back(dim);
       }
     }