[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

5 files changed, 220 insertions, 41 deletions

diff --git a/include/operator.def b/include/operator.def
index 6198c0e..5b0cc7f 100644
--- a/include/operator.def
+++ b/include/operator.def
@@ -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
@@ -119,3 +119,11 @@ DEF_OPERATOR(Custom)
 /* control flow operators */
 DEF_OPERATOR(If)
 DEF_OPERATOR(While)
+
+/* shape operators */
+DEF_OPERATOR(ConstShape)
+DEF_OPERATOR(ConcatShape)
+DEF_OPERATOR(AddShape)
+DEF_OPERATOR(SubShape)
+DEF_OPERATOR(MulShape)
+DEF_OPERATOR(DivShape)
diff --git a/include/schema_operator.def b/include/schema_operator.def
index 52c7ae4..675b4f2 100644
--- a/include/schema_operator.def
+++ b/include/schema_operator.def
@@ -1,4 +1,4 @@
-// Copyright (c) 2023, ARM Limited.
+// Copyright (c) 2023-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,3 +93,9 @@ DEF_SCHEMA_OPERATOR(COND_IF)
 DEF_SCHEMA_OPERATOR(WHILE_LOOP)
 DEF_SCHEMA_OPERATOR(FFT2D)
 DEF_SCHEMA_OPERATOR(RFFT2D)
+DEF_SCHEMA_OPERATOR(CONST_SHAPE)
+DEF_SCHEMA_OPERATOR(CONCAT_SHAPE)
+DEF_SCHEMA_OPERATOR(ADD_SHAPE)
+DEF_SCHEMA_OPERATOR(SUB_SHAPE)
+DEF_SCHEMA_OPERATOR(MUL_SHAPE)
+DEF_SCHEMA_OPERATOR(DIV_SHAPE)
diff --git a/src/TosaDeserialize.cpp b/src/TosaDeserialize.cpp
index f1b7d98..2421d79 100644
--- a/src/TosaDeserialize.cpp
+++ b/src/TosaDeserialize.cpp
@@ -1,5 +1,5 @@
 
-// Copyright (c) 2023, ARM Limited.
+// Copyright (c) 2023-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
@@ -139,8 +139,9 @@ mlir::LogicalResult BuildTensorType(mlir::OpBuilder *op_builder,
     element_type = op_builder->getBF16Type();
     break;
   case DType_SHAPE:
-    element_type = op_builder->getIntegerType(64);
-    break;
+    llvm::errs()
+        << "ERROR: Cannot construct RankedTensorType out of tosa.shape type \n";
+    return mlir::failure();
   default:
     llvm::errs() << "ERROR: unknown type " << EnumNamesDType()[ts->GetDtype()]
                  << "\n";
@@ -236,7 +237,6 @@ ConstructConstAttr(const mlir::RankedTensorType &output_type,
   case DType_UINT8:
   case DType_UINT16:
   case DType_BF16:
-  case DType_SHAPE:
   default: {
     llvm::errs() << "ERROR: " << op_name
                  << " contains unsupported element type\n";
@@ -375,10 +375,11 @@ public:
       mlir::Block *_block, mlir::Location _loc,
       TosaMlirBlockBuilder *_block_builder,
       std::unordered_map<std::string, mlir::Value> *_tensor_map,
-      std::unordered_map<std::string, mlir::RankedTensorType> *_tensor_type_map)
+      std::unordered_map<std::string, mlir::RankedTensorType> *_tensor_type_map,
+      std::unordered_map<std::string, mlir::tosa::shapeType> *_shape_type_map)
       : op_builder(_op_builder), ser_block(_ser_block), block(_block),
         loc(_loc), block_builder(_block_builder), tensor_map(_tensor_map),
-        tensor_type_map(_tensor_type_map) {}
+        tensor_type_map(_tensor_type_map), shape_type_map(_shape_type_map) {}
 
   template <Op OPCODE>
   std::vector<mlir::Value> build(TosaSerializationOperator *op) const;
@@ -420,8 +421,16 @@ private:
 
   template <class MLIR_OP>
   std::vector<mlir::Value>
+  BuildEwiseBinaryShapeOp(TosaSerializationOperator *op) const;
+
+  template <class MLIR_OP>
+  std::vector<mlir::Value>
   BuildReductionOp(TosaSerializationOperator *op) const;
 
+  template <class T>
+  mlir::Value BuildConstShape(mlir::OpBuilder *op_builder, mlir::Location loc,
+                              const std::vector<T> &values) const;
+
   template <class MLIR_OP>
   std::vector<mlir::Value> BuildConvOp(TosaSerializationOperator *op) const;
 
@@ -432,6 +441,7 @@ private:
   TosaMlirBlockBuilder *block_builder;
   std::unordered_map<std::string, mlir::Value> *tensor_map;
   std::unordered_map<std::string, mlir::RankedTensorType> *tensor_type_map;
+  std::unordered_map<std::string, mlir::tosa::shapeType> *shape_type_map;
 };
 
 // Main template to catch unimplemented translation
@@ -566,6 +576,22 @@ std::vector<mlir::Value> TosaMlirOperatorBuilder::BuildEwiseBinaryOp(
 }
 
 template <class MLIR_OP>
+std::vector<mlir::Value> TosaMlirOperatorBuilder::BuildEwiseBinaryShapeOp(
+    TosaSerializationOperator *op) const {
+  mlir::Value input0_val = tensor_map->at(op->GetInputTensorNames()[0]);
+  mlir::Value input1_val = tensor_map->at(op->GetInputTensorNames()[1]);
+  mlir::tosa::shapeType output_type =
+      shape_type_map->at(op->GetOutputTensorNames()[0]);
+  assert(op->GetAttributeType() ==
+         Attribute_NONE); // double check that there is no attribute
+
+  mlir::Operation *mlir_op =
+      op_builder->create<MLIR_OP>(loc, output_type, input0_val, input1_val);
+  block->push_back(mlir_op);
+  return std::vector<mlir::Value>({mlir_op->getResult(0)});
+}
+
+template <class MLIR_OP>
 std::vector<mlir::Value>
 TosaMlirOperatorBuilder::BuildReductionOp(TosaSerializationOperator *op) const {
   mlir::Value input_val = tensor_map->at(op->GetInputTensorNames()[0]);
@@ -600,6 +626,14 @@ TosaMlirOperatorBuilder::BuildReductionOp(TosaSerializationOperator *op) const {
     return BuildEwiseBinaryOp<mlir::tosa::MLIR_OP_NAME##Op>(op);               \
   }
 
+#define BUILD_OP_ELEMENTWISE_BINARY_SHAPE(MLIR_OP_NAME, SCHEMA_OP_NAME)        \
+  template <>                                                                  \
+  std::vector<mlir::Value>                                                     \
+  TosaMlirOperatorBuilder::build<Op_##SCHEMA_OP_NAME>(                         \
+      TosaSerializationOperator * op) const {                                  \
+    return BuildEwiseBinaryShapeOp<mlir::tosa::MLIR_OP_NAME##Op>(op);          \
+  }
+
 #define BUILD_OP_REDUCTION(MLIR_OP_NAME, SCHEMA_OP_NAME)                       \
   template <>                                                                  \
   std::vector<mlir::Value>                                                     \
@@ -654,6 +688,11 @@ BUILD_OP_ELEMENTWISE_UNARY(Tanh, TANH)
 BUILD_OP_ELEMENTWISE_UNARY(Identity, IDENTITY)
 BUILD_OP_ELEMENTWISE_UNARY(Cast, CAST)
 
+BUILD_OP_ELEMENTWISE_BINARY_SHAPE(AddShape, ADD_SHAPE)
+BUILD_OP_ELEMENTWISE_BINARY_SHAPE(SubShape, SUB_SHAPE)
+BUILD_OP_ELEMENTWISE_BINARY_SHAPE(MulShape, MUL_SHAPE)
+BUILD_OP_ELEMENTWISE_BINARY_SHAPE(DivShape, DIV_SHAPE)
+
 template <>
 std::vector<mlir::Value>
 TosaMlirOperatorBuilder::build<Op_CONST>(TosaSerializationOperator *op) const {
@@ -670,6 +709,40 @@ TosaMlirOperatorBuilder::build<Op_CONST>(TosaSerializationOperator *op) const {
   return std::vector<mlir::Value>({mlir_op->getResult(0)});
 }
 
+template <class T>
+mlir::Value
+TosaMlirOperatorBuilder::BuildConstShape(mlir::OpBuilder *op_builder,
+                                         mlir::Location loc,
+                                         const std::vector<T> &values) const {
+  std::vector<int64_t> vec;
+  for (auto val : values) {
+    vec.push_back(val);
+  }
+  auto attr = op_builder->getIndexTensorAttr(vec);
+  auto type = mlir::tosa::shapeType::get(op_builder->getContext(),
+                                         /* rank = */ vec.size());
+  mlir::Operation *mlir_op =
+      op_builder->create<mlir::tosa::ConstShapeOp>(loc, type, attr);
+  block->push_back(mlir_op);
+  return mlir_op->getResult(0);
+}
+
+template <>
+std::vector<mlir::Value> TosaMlirOperatorBuilder::build<Op_CONST_SHAPE>(
+    TosaSerializationOperator *op) const {
+  const auto &output_name = op->GetOutputTensorNames()[0];
+  mlir::tosa::shapeType output_type = shape_type_map->at(output_name);
+  TosaSerializationTensor *ts = ser_block->GetTensorByName(output_name);
+
+  const auto &data = ts->GetData();
+
+  std::vector<int64_t> i64_data;
+  TosaSerializationHandler::ConvertU8toI64(data, output_type.getRank(),
+                                           i64_data);
+  mlir::Value result = BuildConstShape(op_builder, loc, i64_data);
+  return std::vector<mlir::Value>({result});
+}
+
 template <class MLIR_OP>
 std::vector<mlir::Value>
 TosaMlirOperatorBuilder::BuildConvOp(TosaSerializationOperator *op) const {
@@ -891,6 +964,24 @@ TosaMlirOperatorBuilder::build<Op_CONCAT>(TosaSerializationOperator *op) const {
 }
 
 template <>
+std::vector<mlir::Value> TosaMlirOperatorBuilder::build<Op_CONCAT_SHAPE>(
+    TosaSerializationOperator *op) const {
+  mlir::tosa::shapeType output_type =
+      shape_type_map->at(op->GetOutputTensorNames()[0]);
+
+  llvm::SmallVector<mlir::Value> input_values;
+  for (auto &input_name : op->GetInputTensorNames()) {
+    mlir::Value input_val = tensor_map->at(input_name);
+    input_values.push_back(input_val);
+  }
+
+  mlir::Operation *mlir_op = op_builder->create<mlir::tosa::ConcatShapeOp>(
+      loc, output_type, input_values);
+  block->push_back(mlir_op);
+  return std::vector<mlir::Value>({mlir_op->getResult(0)});
+}
+
+template <>
 std::vector<mlir::Value>
 TosaMlirOperatorBuilder::build<Op_NEGATE>(TosaSerializationOperator *op) const {
   mlir::Value input_val = tensor_map->at(op->GetInputTensorNames()[0]);
@@ -926,17 +1017,10 @@ std::vector<mlir::Value> TosaMlirOperatorBuilder::build<Op_RESHAPE>(
   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_ReshapeAttribute); // double check attribute type
-  TosaReshapeAttribute *attr =
-      static_cast<TosaReshapeAttribute *>(op->GetAttribute());
-
-  mlir::DenseI64ArrayAttr new_shape =
-      BuildDenseI64ArrayAttr(op_builder, attr->new_shape());
+  mlir::Value shape_val = tensor_map->at(op->GetInputTensorNames()[1]);
 
   mlir::Operation *mlir_op = op_builder->create<mlir::tosa::ReshapeOp>(
-      loc, output_type, input_val, new_shape);
+      loc, output_type, input_val, shape_val);
   block->push_back(mlir_op);
   return std::vector<mlir::Value>({mlir_op->getResult(0)});
 }
@@ -1081,16 +1165,12 @@ template <>
 std::vector<mlir::Value>
 TosaMlirOperatorBuilder::build<Op_TILE>(TosaSerializationOperator *op) const {
   mlir::Value input_val = tensor_map->at(op->GetInputTensorNames()[0]);
+  mlir::Value multiples = tensor_map->at(op->GetInputTensorNames()[1]);
   mlir::RankedTensorType output_type =
       tensor_type_map->at(op->GetOutputTensorNames()[0]);
 
   assert(op->GetAttributeType() ==
-         Attribute_TileAttribute); // double check attribute type
-  TosaTileAttribute *attr =
-      static_cast<TosaTileAttribute *>(op->GetAttribute());
-
-  mlir::DenseI64ArrayAttr multiples =
-      BuildDenseI64ArrayAttr(op_builder, attr->multiples());
+         Attribute_NONE); // double check attribute type
 
   mlir::Operation *mlir_op = op_builder->create<mlir::tosa::TileOp>(
       loc, output_type, input_val, multiples);
@@ -1395,6 +1475,7 @@ private:
   TosaMlirRegionBuilder *region_builder;
   mlir::Block *block;
   std::unordered_map<std::string, mlir::RankedTensorType> tensor_type_map;
+  std::unordered_map<std::string, mlir::tosa::shapeType> shape_type_map;
   std::unordered_set<std::string> unranked_tensors;
 };
 
@@ -1562,17 +1643,25 @@ mlir::LogicalResult TosaMlirBlockBuilder::BuildAllOpsInBlock(
   std::queue<TosaSerializationOperator *> operator_queue;
 
   TosaMlirOperatorBuilder tosa_op_builder(op_builder, ser_block, block, loc,
-                                          this, &tensor_map, &tensor_type_map);
+                                          this, &tensor_map, &tensor_type_map,
+                                          &shape_type_map);
 
   for (auto ts : ser_block->GetTensors()) {
     if (ts->GetVariable()) {
       RegisterVariableTensor(ts);
     }
+    const auto &ts_name = ts->GetName();
+    if (ts->GetDtype() == DType::DType_SHAPE) {
+      // ts is tosa.shape type
+      auto shape_rank = ts->GetShape()[0];
+      shape_type_map[ts_name] =
+          mlir::tosa::shapeType::get(op_builder->getContext(), shape_rank);
+      continue;
+    }
     mlir::RankedTensorType type;
     if (BuildTensorType(op_builder, ts, type).failed()) {
       return mlir::failure();
     }
-    const auto &ts_name = ts->GetName();
     tensor_type_map[ts_name] = type;
     if (ts->GetIsUnranked()) {
       assert(ts->GetShape().empty()); // unranked tensors should have shape = {}
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);
       }
     }
diff --git a/third_party/serialization_lib b/third_party/serialization_lib
-Subproject cd914da80304592a6826e967d0d04b637d91671
+Subproject 5d580faec02bcef56164587accb5fd88a3e80d8