[tosa_mlir_translator] Add local_bound support

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

3 files changed, 66 insertions, 27 deletions

diff --git a/src/TosaDeserialize.cpp b/src/TosaDeserialize.cpp
index 21798f3..9649644 100644
--- a/src/TosaDeserialize.cpp
+++ b/src/TosaDeserialize.cpp
@@ -673,22 +673,21 @@ TosaMlirOperatorBuilder::BuildConvOp(TosaSerializationOperator *op) const {
       BuildDenseI64ArrayAttr(op_builder, attr->dilation());
   auto input_zp = attr->input_zp();
   auto weight_zp = attr->weight_zp();
+  bool local_bound = attr->local_bound();
 
   // quantizationattr is required for quantized type, and not allowed for float
   // type
   mlir::Operation *mlir_op;
   if (output_type.getElementType().isa<mlir::FloatType>()) {
     assert(input_zp == 0 && weight_zp == 0);
-    mlir_op =
-        op_builder->create<MLIR_OP>(loc, output_type, input0_val, input1_val,
-                                    input2_val, pad, stride, dilation);
-  } else {
-    auto quant = op_builder->getAttr<mlir::tosa::ConvOpQuantizationAttr>(
-        input_zp, weight_zp);
-    mlir_op =
-        op_builder->create<MLIR_OP>(loc, output_type, input0_val, input1_val,
-                                    input2_val, pad, stride, dilation, quant);
   }
+
+  auto quant = op_builder->getAttr<mlir::tosa::ConvOpQuantizationAttr>(
+      input_zp, weight_zp);
+  mlir_op = op_builder->create<MLIR_OP>(loc, output_type, input0_val,
+                                        input1_val, input2_val, pad, stride,
+                                        dilation, quant, local_bound);
+
   block->push_back(mlir_op);
   return std::vector<mlir::Value>({mlir_op->getResult(0)});
 }
@@ -726,22 +725,20 @@ std::vector<mlir::Value> TosaMlirOperatorBuilder::build<Op_TRANSPOSE_CONV2D>(
       BuildDenseI64ArrayAttr(op_builder, attr->output_shape());
   auto input_zp = attr->input_zp();
   auto weight_zp = attr->weight_zp();
+  bool local_bound = attr->local_bound();
 
   // quantizationattr is required for quantized type, and not allowed for float
   // type
   mlir::Operation *mlir_op;
   if (output_type.getElementType().isa<mlir::FloatType>()) {
     assert(input_zp == 0 && weight_zp == 0);
-    mlir_op = op_builder->create<mlir::tosa::TransposeConv2DOp>(
-        loc, output_type, input0_val, input1_val, input2_val, out_pad, stride,
-        output_shape);
-  } else {
-    auto quant = op_builder->getAttr<mlir::tosa::ConvOpQuantizationAttr>(
-        input_zp, weight_zp);
-    mlir_op = op_builder->create<mlir::tosa::TransposeConv2DOp>(
-        loc, output_type, input0_val, input1_val, input2_val, out_pad, stride,
-        output_shape, quant);
   }
+  auto quant = op_builder->getAttr<mlir::tosa::ConvOpQuantizationAttr>(
+      input_zp, weight_zp);
+  mlir_op = op_builder->create<mlir::tosa::TransposeConv2DOp>(
+      loc, output_type, input0_val, input1_val, input2_val, out_pad, stride,
+      output_shape, quant, local_bound);
+
   block->push_back(mlir_op);
   return std::vector<mlir::Value>({mlir_op->getResult(0)});
 }
@@ -1283,10 +1280,14 @@ TosaMlirOperatorBuilder::build<Op_RFFT2D>(TosaSerializationOperator *op) const {
   mlir::RankedTensorType output1_type =
       tensor_type_map->at(op->GetOutputTensorNames()[1]);
   assert(op->GetAttributeType() ==
-         Attribute_NONE); // double check that there is no attribute
+         Attribute_RFFTAttribute); // double check attribute type
+  TosaRFFTAttribute *attr =
+      static_cast<TosaRFFTAttribute *>(op->GetAttribute());
+
+  bool local_bound = attr->local_bound();
 
   mlir::Operation *mlir_op = op_builder->create<mlir::tosa::RFFT2dOp>(
-      loc, output0_type, output1_type, input_val);
+      loc, output0_type, output1_type, input_val, local_bound);
   block->push_back(mlir_op);
   return std::vector<mlir::Value>(
       {mlir_op->getResult(0), mlir_op->getResult(1)});
@@ -1305,9 +1306,11 @@ TosaMlirOperatorBuilder::build<Op_FFT2D>(TosaSerializationOperator *op) const {
   assert(op->GetAttributeType() == Attribute_FFTAttribute);
   TosaFFTAttribute *attr = static_cast<TosaFFTAttribute *>(op->GetAttribute());
   auto inverse = op_builder->getBoolAttr(attr->inverse());
+  auto local_bound = op_builder->getBoolAttr(attr->local_bound());
 
   mlir::Operation *mlir_op = op_builder->create<mlir::tosa::FFT2dOp>(
-      loc, output0_type, output1_type, input0_val, input1_val, inverse);
+      loc, output0_type, output1_type, input0_val, input1_val, inverse,
+      local_bound);
   block->push_back(mlir_op);
   return std::vector<mlir::Value>(
       {mlir_op->getResult(0), mlir_op->getResult(1)});
diff --git a/src/TosaSerialize.cpp b/src/TosaSerialize.cpp
index 8aef8fd..a131117 100644
--- a/src/TosaSerialize.cpp
+++ b/src/TosaSerialize.cpp
@@ -790,7 +790,13 @@ TosaSerializationOperatorBuilder::build<mlir::tosa::Conv2DOp>(
   mlir::RankedTensorType tensor =
       op.getOperand(0).getType().cast<mlir::RankedTensorType>();
 
-  TosaConvAttribute attribute(pad, stride, dilation, input_zp, weight_zp);
+  bool local_bound =
+      op.hasAttr("local_bound")
+          ? op.getAttr("local_bound").dyn_cast<mlir::BoolAttr>().getValue()
+          : false;
+
+  TosaConvAttribute attribute(pad, stride, dilation, input_zp, weight_zp,
+                              local_bound);
 
   TosaSerializationOperator *tyop = new TosaSerializationOperator(
       Op_CONV2D, Attribute_ConvAttribute, &attribute,
@@ -826,7 +832,13 @@ TosaSerializationOperatorBuilder::build<mlir::tosa::Conv3DOp>(
   mlir::RankedTensorType tensor =
       op.getOperand(0).getType().cast<mlir::RankedTensorType>();
 
-  TosaConvAttribute attribute(pad, stride, dilation, input_zp, weight_zp);
+  bool local_bound =
+      op.hasAttr("local_bound")
+          ? op.getAttr("local_bound").dyn_cast<mlir::BoolAttr>().getValue()
+          : false;
+
+  TosaConvAttribute attribute(pad, stride, dilation, input_zp, weight_zp,
+                              local_bound);
 
   TosaSerializationOperator *tyop = new TosaSerializationOperator(
       Op_CONV3D, Attribute_ConvAttribute, &attribute,
@@ -862,7 +874,13 @@ TosaSerializationOperatorBuilder::build<mlir::tosa::DepthwiseConv2DOp>(
   mlir::RankedTensorType tensor =
       op.getOperand(0).getType().cast<mlir::RankedTensorType>();
 
-  TosaConvAttribute attribute(pad, stride, dilation, input_zp, weight_zp);
+  bool local_bound =
+      op.hasAttr("local_bound")
+          ? op.getAttr("local_bound").dyn_cast<mlir::BoolAttr>().getValue()
+          : false;
+
+  TosaConvAttribute attribute(pad, stride, dilation, input_zp, weight_zp,
+                              local_bound);
 
   TosaSerializationOperator *tyop = new TosaSerializationOperator(
       Op_DEPTHWISE_CONV2D, Attribute_ConvAttribute, &attribute,
@@ -898,8 +916,13 @@ TosaSerializationOperatorBuilder::build<mlir::tosa::TransposeConv2DOp>(
   mlir::RankedTensorType tensor =
       op.getOperand(0).getType().cast<mlir::RankedTensorType>();
 
+  bool local_bound =
+      op.hasAttr("local_bound")
+          ? op.getAttr("local_bound").dyn_cast<mlir::BoolAttr>().getValue()
+          : false;
+
   TosaTransposeConvAttribute attribute(out_pad, stride, out_shape, input_zp,
-                                       weight_zp);
+                                       weight_zp, local_bound);
 
   TosaSerializationOperator *tyop = new TosaSerializationOperator(
       Op_TRANSPOSE_CONV2D, Attribute_TransposeConvAttribute, &attribute,
@@ -1597,8 +1620,16 @@ TosaSerializationOperatorBuilder::build<mlir::tosa::RFFT2dOp>(
   std::string output_real_name = GetTensorName(op.getResult(0));
   std::string output_imag_name = GetTensorName(op.getResult(1));
 
+  bool local_bound =
+      op.hasAttr("local_bound")
+          ? op.getAttr("local_bound").dyn_cast<mlir::BoolAttr>().getValue()
+          : false;
+
+  TosaRFFTAttribute attribute(local_bound);
+
   TosaSerializationOperator *tyop = new TosaSerializationOperator(
-      Op_RFFT2D, Attribute_NONE, nullptr, std::vector<std::string>{input_name},
+      Op_RFFT2D, Attribute_RFFTAttribute, &attribute,
+      std::vector<std::string>{input_name},
       std::vector<std::string>{output_real_name, output_imag_name});
 
   return tyop;
@@ -1611,12 +1642,17 @@ TosaSerializationOperatorBuilder::build<mlir::tosa::FFT2dOp>(
 
   bool inverse = op.getAttr("inverse").dyn_cast<mlir::BoolAttr>().getValue();
 
+  bool local_bound =
+      op.hasAttr("local_bound")
+          ? op.getAttr("local_bound").dyn_cast<mlir::BoolAttr>().getValue()
+          : false;
+
   std::string input_real_name = GetTensorName(op.getOperand(0));
   std::string input_imag_name = GetTensorName(op.getOperand(1));
   std::string output_real_name = GetTensorName(op.getResult(0));
   std::string output_imag_name = GetTensorName(op.getResult(1));
 
-  TosaFFTAttribute attribute(inverse);
+  TosaFFTAttribute attribute(inverse, local_bound);
 
   TosaSerializationOperator *tyop = new TosaSerializationOperator(
       Op_FFT2D, Attribute_FFTAttribute, &attribute,
diff --git a/third_party/serialization_lib b/third_party/serialization_lib
-Subproject 5917fc7a9392da8fd1e8c68b2d00b89709a3158
+Subproject 4881c29247d4b411de446b13d9bd58ea93737aa