1 files changed, 112 insertions, 165 deletions

diff --git a/src/TosaDeserialize.cpp b/src/TosaDeserialize.cpp
index 6fa691e..fdbd892 100644
--- a/src/TosaDeserialize.cpp
+++ b/src/TosaDeserialize.cpp
@@ -166,58 +166,70 @@ mlir::LogicalResult BuildTensorType(mlir::OpBuilder *op_builder,
   return mlir::success();
 }
 
-mlir::DenseElementsAttr
-ConstructConstAttr(const mlir::RankedTensorType &output_type,
-                   TosaSerializationTensor *ts, const std::string &op_name) {
-  const auto &data = ts->GetData();
-  auto &shape = ts->GetShape();
-  // compute output data size
-  uint32_t out_size = 1;
-  for (const auto dim : shape) {
-    out_size *= dim;
-  }
-  mlir::DenseElementsAttr value_attr;
-  switch (ts->GetDtype()) {
-  case DType_FP32:
-  case DType_BF16:
-  case DType_FP8E4M3:
-  case DType_FP8E5M2: {
-    // for FP32, FP16 and FP8 types, value attributes are stored as FP32 values
+mlir::DenseElementsAttr GetConstAttr(const std::vector<uint8_t> &data,
+                                     const mlir::RankedTensorType &output_type,
+                                     uint32_t out_size) {
+  auto element_type = output_type.getElementType();
+  if (element_type.isF32()) {
+    // for FP32, value attributes are stored as FP32 values
     std::vector<float> float_data;
     TosaSerializationHandler::ConvertU8toF32(data, out_size, float_data);
-    value_attr =
-        mlir::DenseElementsAttr::get(output_type, llvm::ArrayRef(float_data));
-    break;
+    return mlir::DenseElementsAttr::get(output_type,
+                                        llvm::ArrayRef(float_data));
+  }
+  if (element_type.isBF16()) {
+    mlir::SmallVector<mlir::APFloat> bf16_data;
+    for (uint32_t i = 0; i < out_size; i++) {
+      uint64_t byte0 = data[i * sizeof(int16_t)];
+      uint64_t byte1 = data[i * sizeof(int16_t) + 1];
+      uint64_t bits = byte0 + (byte1 << 8);
+      mlir::APInt bf16_bits(16, bits);
+      mlir::APFloat bf16(mlir::APFloat::BFloat(), bf16_bits);
+      bf16_data.push_back(bf16);
+    }
+    return mlir::DenseElementsAttr::get(output_type, bf16_data);
+  }
+  if (element_type.isFloat8E4M3FN()) {
+    mlir::SmallVector<mlir::APFloat> f8_data;
+    for (uint32_t i = 0; i < out_size; i++) {
+      mlir::APInt f8_bits(8, static_cast<uint64_t>(data[i]));
+      mlir::APFloat f8(mlir::APFloat::Float8E4M3FN(), f8_bits);
+      f8_data.push_back(f8);
+    }
+    return mlir::DenseElementsAttr::get(output_type, f8_data);
+  }
+  if (element_type.isFloat8E5M2()) {
+    mlir::SmallVector<mlir::APFloat> f8_data;
+    for (uint32_t i = 0; i < out_size; i++) {
+      mlir::APInt f8_bits(8, static_cast<uint64_t>(data[i]));
+      mlir::APFloat f8(mlir::APFloat::Float8E5M2(), f8_bits);
+      f8_data.push_back(f8);
+    }
+    return mlir::DenseElementsAttr::get(output_type, f8_data);
   }
-  case DType_INT4: {
+  if (element_type.isInteger(4)) {
     std::vector<int8_t> int4_data;
     TosaSerializationHandler::ConvertU8toI4(data, out_size, int4_data);
-    value_attr =
-        mlir::DenseElementsAttr::get(output_type, llvm::ArrayRef(int4_data));
-    break;
+    return mlir::DenseElementsAttr::get(output_type, llvm::ArrayRef(int4_data));
   }
-  case DType_INT8: {
+  if (element_type.isInteger(8)) {
     std::vector<int8_t> int8_data;
     TosaSerializationHandler::ConvertU8toI8(data, out_size, int8_data);
-    value_attr =
-        mlir::DenseElementsAttr::get(output_type, llvm::ArrayRef(int8_data));
-    break;
+    return mlir::DenseElementsAttr::get(output_type, llvm::ArrayRef(int8_data));
   }
-  case DType_INT16: {
+  if (element_type.isInteger(16)) {
     std::vector<int16_t> int16_data;
     TosaSerializationHandler::ConvertU8toI16(data, out_size, int16_data);
-    value_attr =
-        mlir::DenseElementsAttr::get(output_type, llvm::ArrayRef(int16_data));
-    break;
+    return mlir::DenseElementsAttr::get(output_type,
+                                        llvm::ArrayRef(int16_data));
   }
-  case DType_INT32: {
+  if (element_type.isInteger(32)) {
     std::vector<int32_t> int32_data;
     TosaSerializationHandler::ConvertU8toI32(data, out_size, int32_data);
-    value_attr =
-        mlir::DenseElementsAttr::get(output_type, llvm::ArrayRef(int32_data));
-    break;
+    return mlir::DenseElementsAttr::get(output_type,
+                                        llvm::ArrayRef(int32_data));
   }
-  case DType_INT48: {
+  if (element_type.isInteger(48)) {
     std::vector<int64_t> int48_data;
     TosaSerializationHandler::ConvertU8toI48(data, out_size, int48_data);
     std::vector<mlir::APInt> apint_data;
@@ -226,34 +238,38 @@ ConstructConstAttr(const mlir::RankedTensorType &output_type,
                               /* isSigned = */ false);
       apint_data.push_back(apint_value);
     }
-    value_attr =
-        mlir::DenseElementsAttr::get(output_type, llvm::ArrayRef(apint_data));
-    break;
+    return mlir::DenseElementsAttr::get(output_type,
+                                        llvm::ArrayRef(apint_data));
   }
-  case DType_BOOL: {
+  if (element_type.isInteger(1)) {
     std::vector<bool> bool_data;
     TosaSerializationHandler::ConvertU8toBool(data, out_size, bool_data);
     llvm::SmallVector<bool> bool_values(bool_data.begin(), bool_data.end());
-    value_attr = mlir::DenseElementsAttr::get(output_type, bool_values);
-    break;
+    return mlir::DenseElementsAttr::get(output_type, bool_values);
   }
-  case DType_FP16: {
+  if (element_type.isF16()) {
     std::vector<half_float::half> half_data;
     TosaSerializationHandler::ConvertU8toF16(data, out_size, half_data);
-    value_attr =
-        mlir::DenseElementsAttr::get(output_type, llvm::ArrayRef(half_data));
-    break;
+    return mlir::DenseElementsAttr::get(output_type, llvm::ArrayRef(half_data));
   }
-  case DType_UINT8:
-  case DType_UINT16:
-  default: {
+
+  return nullptr;
+}
+
+mlir::DenseElementsAttr
+ConstructConstAttr(const mlir::RankedTensorType &output_type,
+                   TosaSerializationTensor *ts, const std::string &op_name) {
+  // compute output data size
+  uint32_t out_size = 1;
+  for (const auto dim : ts->GetShape()) {
+    out_size *= dim;
+  }
+  auto attr = GetConstAttr(ts->GetData(), output_type, out_size);
+  if (!attr) {
     llvm::errs() << "ERROR: " << op_name
                  << " contains unsupported element type\n";
-    return nullptr;
   }
-  }
-
-  return value_attr;
+  return attr;
 }
 
 mlir::LogicalResult ConstructVariableOps(mlir::ModuleOp &module) {
@@ -942,56 +958,31 @@ TosaMlirOperatorBuilder::build<Op_CLAMP>(TosaSerializationOperator *op) const {
   TosaClampAttribute *attr =
       static_cast<TosaClampAttribute *>(op->GetAttribute());
 
-  mlir::Type input_element_type =
+  mlir::Type element_type =
       llvm::cast<mlir::ShapedType>(input_val.getType()).getElementType();
-  if (auto quantType = llvm::dyn_cast<mlir::quant::UniformQuantizedType>(
-          input_element_type)) {
-    input_element_type = quantType.getStorageType();
+  if (auto quantType =
+          llvm::dyn_cast<mlir::quant::UniformQuantizedType>(element_type)) {
+    element_type = quantType.getStorageType();
   }
 
+  auto element_const_type = mlir::RankedTensorType::get({1}, element_type);
+  auto min_values_attr = GetConstAttr(attr->min_val(), element_const_type, 1);
+  auto max_values_attr = GetConstAttr(attr->max_val(), element_const_type, 1);
+
   mlir::Attribute min_val_attr, max_val_attr;
-  if (input_element_type.isa<mlir::FloatType>()) {
-    std::vector<float> min_float_data, max_float_data;
-    TosaSerializationHandler::ConvertU8toF32(attr->min_val(), /* size = */ 1,
-                                             min_float_data);
-    TosaSerializationHandler::ConvertU8toF32(attr->max_val(), /* size = */ 1,
-                                             max_float_data);
-    min_val_attr =
-        op_builder->getFloatAttr(input_element_type, min_float_data[0]);
-    max_val_attr =
-        op_builder->getFloatAttr(input_element_type, max_float_data[0]);
+  if (element_type.isa<mlir::FloatType>()) {
+    min_val_attr = op_builder->getFloatAttr(
+        element_type, min_values_attr.getValues<mlir::APFloat>()[0]);
+    max_val_attr = op_builder->getFloatAttr(
+        element_type, max_values_attr.getValues<mlir::APFloat>()[0]);
   } else {
-    std::vector<int32_t> min_int_data, max_int_data;
-    TosaSerializationHandler::ConvertU8toI32(attr->min_val(), /* size = */ 1,
-                                             min_int_data);
-    TosaSerializationHandler::ConvertU8toI32(attr->max_val(), /* size = */ 1,
-                                             max_int_data);
-    if (input_element_type.isUnsignedInteger()) {
-      if (input_element_type.isUnsignedInteger(8)) {
-        uint8_t min_val = min_int_data[0];
-        uint8_t max_val = max_int_data[0];
-        min_val_attr = op_builder->getIntegerAttr(input_element_type, min_val);
-        max_val_attr = op_builder->getIntegerAttr(input_element_type, max_val);
-      } else if (input_element_type.isUnsignedInteger(16)) {
-        uint16_t min_val = min_int_data[0];
-        uint16_t max_val = max_int_data[0];
-        min_val_attr = op_builder->getIntegerAttr(input_element_type, min_val);
-        max_val_attr = op_builder->getIntegerAttr(input_element_type, max_val);
-      } else {
-        llvm::errs()
-            << "ERROR: " << get_string(op)
-            << " contains unsupported unsigned int element data type.\n";
-        return {};
-      }
-    } else {
-      min_val_attr =
-          op_builder->getIntegerAttr(input_element_type, min_int_data[0]);
-      max_val_attr =
-          op_builder->getIntegerAttr(input_element_type, max_int_data[0]);
-    }
+    min_val_attr = op_builder->getIntegerAttr(
+        element_type, min_values_attr.getValues<mlir::APInt>()[0]);
+    max_val_attr = op_builder->getIntegerAttr(
+        element_type, max_values_attr.getValues<mlir::APInt>()[0]);
   }
 
-  mlir::Operation *mlir_op = op_builder->create<mlir::tosa::ClampOp>(
+  auto mlir_op = op_builder->create<mlir::tosa::ClampOp>(
       loc, output_type, input_val, min_val_attr, max_val_attr);
   block->push_back(mlir_op);
   return std::vector<mlir::Value>({mlir_op->getResult(0)});
@@ -1101,80 +1092,36 @@ TosaMlirOperatorBuilder::build<Op_PAD>(TosaSerializationOperator *op) const {
   assert(op->GetAttributeType() ==
          Attribute_PadAttribute); // double check attribute type
   TosaPadAttribute *attr = static_cast<TosaPadAttribute *>(op->GetAttribute());
-
-  float pad_const_fp = 0.0f;
-  int32_t pad_const_int = 0;
-
-  if (element_type.isa<mlir::FloatType>()) {
-    std::vector<float> float_data;
-    TosaSerializationHandler::ConvertU8toF32(attr->pad_const(),
-                                             /* size = */ 1, float_data);
-    pad_const_fp = float_data[0];
-  } else {
-    std::vector<int32_t> int32_data;
-    TosaSerializationHandler::ConvertU8toI32(attr->pad_const(),
-                                             /* size = */ 1, int32_data);
-    pad_const_int = int32_data[0];
+  const auto &pad_const_u8_data = attr->pad_const();
+
+  // check for any value in pad_const_u8_data
+  bool has_pad_const = false;
+  for (auto v : pad_const_u8_data) {
+    if (v != 0) {
+      has_pad_const = true;
+      break;
+    }
   }
-
-  // todo: int input_zp = attr->pad_input_zp();
-
-  mlir::Operation *mlir_op;
-  mlir::Value pad_const_value;
-
-  bool isBoolType = element_type.isInteger(1);
-  // First handle boolean type.
-  if (isBoolType) {
-    mlir::Type boolType = op_builder->getIntegerType(1);
-    auto pad_const_type = mlir::RankedTensorType::get({}, boolType);
-    // Treat zero integer is `false`, and any non-zero integner evaluates to
-    // `true`.
-    bool pad_const = pad_const_int == 0 ? false : true;
-    auto pad_const_attr =
-        mlir::DenseElementsAttr::get(pad_const_type, {pad_const});
-    mlir::Operation *pad_const_op = op_builder->create<mlir::tosa::ConstOp>(
-        loc, pad_const_type, pad_const_attr);
-
-    block->push_back(pad_const_op);
-    pad_const_value = pad_const_op->getResult(0);
-    mlir_op = op_builder->create<mlir::tosa::PadOp>(
-        loc, output_type, input_val, padding_val, pad_const_value);
-
+  if (!has_pad_const) {
+    // handle the cases where no explicit pad_const input.
+    auto mlir_op = op_builder->create<mlir::tosa::PadOp>(
+        loc, output_type, input_val, padding_val);
     block->push_back(mlir_op);
     return std::vector<mlir::Value>({mlir_op->getResult(0)});
   }
 
-  // Second handle the cases where no explicit pad_const input.
-  if (pad_const_int == 0 && pad_const_fp == 0.0f) {
-    mlir_op = op_builder->create<mlir::tosa::PadOp>(loc, output_type, input_val,
-                                                    padding_val);
-    block->push_back(mlir_op);
-    return std::vector<mlir::Value>({mlir_op->getResult(0)});
-  }
+  // has pad const - create a const op for pad_const input
+  auto pad_const_type = mlir::RankedTensorType::get({}, element_type);
+  auto pad_const_attr = GetConstAttr(pad_const_u8_data, pad_const_type, 1);
 
-  // Then handle explicit numerical pad_const cases.
-  if (pad_const_int != 0) {
-    assert(pad_const_fp == 0.0f && llvm::isa<IntegerType>(element_type));
-    auto pad_const_int_type = mlir::RankedTensorType::get({}, element_type);
-    auto pad_const_int_attr =
-        mlir::DenseElementsAttr::get(pad_const_int_type, {pad_const_int});
-    mlir::Operation *pad_const_int_op = op_builder->create<mlir::tosa::ConstOp>(
-        loc, pad_const_int_type, pad_const_int_attr);
-    block->push_back(pad_const_int_op);
-    pad_const_value = pad_const_int_op->getResult(0);
-  } else {
-    assert(pad_const_fp != 0 && llvm::isa<FloatType>(element_type));
-    auto pad_const_fp_type = mlir::RankedTensorType::get({}, element_type);
-    auto pad_const_fp_attr =
-        mlir::DenseElementsAttr::get(pad_const_fp_type, {pad_const_fp});
-    mlir::Operation *pad_const_fp_op = op_builder->create<mlir::tosa::ConstOp>(
-        loc, pad_const_fp_type, pad_const_fp_attr);
-    block->push_back(pad_const_fp_op);
-    pad_const_value = pad_const_fp_op->getResult(0);
-  }
-
-  mlir_op = op_builder->create<mlir::tosa::PadOp>(loc, output_type, input_val,
-                                                  padding_val, pad_const_value);
+  auto pad_const_op = op_builder->create<mlir::tosa::ConstOp>(
+      loc, pad_const_type, pad_const_attr);
+
+  block->push_back(pad_const_op);
+  mlir::Value pad_const_value = pad_const_op->getResult(0);
+
+  auto mlir_op = op_builder->create<mlir::tosa::PadOp>(
+      loc, output_type, input_val, padding_val, pad_const_value);
 
   block->push_back(mlir_op);
   return std::vector<mlir::Value>({mlir_op->getResult(0)});