diff options
author | Tai Ly <tai.ly@arm.com> | 2024-01-16 19:52:58 +0000 |
---|---|---|
committer | Tai Ly <tai.ly@arm.com> | 2024-01-18 23:46:37 +0000 |
commit | 4d04238050a50081f35d7e615b590d00e702f501 (patch) | |
tree | 045fbbc352d9b988ca501c3557adf4124879111f | |
parent | 20f6941b21f84cd5f0152d42f343b0992dd5a6e5 (diff) | |
download | tosa_mlir_translator-4d04238050a50081f35d7e615b590d00e702f501.tar.gz |
[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
-rw-r--r-- | include/operator.def | 10 | ||||
-rw-r--r-- | include/schema_operator.def | 8 | ||||
-rw-r--r-- | src/TosaDeserialize.cpp | 135 | ||||
-rw-r--r-- | src/TosaSerialize.cpp | 108 | ||||
m--------- | third_party/serialization_lib | 0 |
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 |