From 4d04238050a50081f35d7e615b590d00e702f501 Mon Sep 17 00:00:00 2001 From: Tai Ly Date: Tue, 16 Jan 2024 19:52:58 +0000 Subject: [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 Change-Id: I4329e621fd7637b1a3491c195fbda77d2a0ad23a --- src/TosaDeserialize.cpp | 135 +++++++++++++++++++++++++++++++++++++++--------- src/TosaSerialize.cpp | 108 ++++++++++++++++++++++++++++++++------ 2 files changed, 204 insertions(+), 39 deletions(-) (limited to 'src') 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 *_tensor_map, - std::unordered_map *_tensor_type_map) + std::unordered_map *_tensor_type_map, + std::unordered_map *_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 std::vector build(TosaSerializationOperator *op) const; @@ -418,10 +419,18 @@ private: std::vector BuildEwiseBinaryOp(TosaSerializationOperator *op) const; + template + std::vector + BuildEwiseBinaryShapeOp(TosaSerializationOperator *op) const; + template std::vector BuildReductionOp(TosaSerializationOperator *op) const; + template + mlir::Value BuildConstShape(mlir::OpBuilder *op_builder, mlir::Location loc, + const std::vector &values) const; + template std::vector BuildConvOp(TosaSerializationOperator *op) const; @@ -432,6 +441,7 @@ private: TosaMlirBlockBuilder *block_builder; std::unordered_map *tensor_map; std::unordered_map *tensor_type_map; + std::unordered_map *shape_type_map; }; // Main template to catch unimplemented translation @@ -565,6 +575,22 @@ std::vector TosaMlirOperatorBuilder::BuildEwiseBinaryOp( return std::vector({mlir_op->getResult(0)}); } +template +std::vector 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(loc, output_type, input0_val, input1_val); + block->push_back(mlir_op); + return std::vector({mlir_op->getResult(0)}); +} + template std::vector TosaMlirOperatorBuilder::BuildReductionOp(TosaSerializationOperator *op) const { @@ -600,6 +626,14 @@ TosaMlirOperatorBuilder::BuildReductionOp(TosaSerializationOperator *op) const { return BuildEwiseBinaryOp(op); \ } +#define BUILD_OP_ELEMENTWISE_BINARY_SHAPE(MLIR_OP_NAME, SCHEMA_OP_NAME) \ + template <> \ + std::vector \ + TosaMlirOperatorBuilder::build( \ + TosaSerializationOperator * op) const { \ + return BuildEwiseBinaryShapeOp(op); \ + } + #define BUILD_OP_REDUCTION(MLIR_OP_NAME, SCHEMA_OP_NAME) \ template <> \ std::vector \ @@ -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 TosaMlirOperatorBuilder::build(TosaSerializationOperator *op) const { @@ -670,6 +709,40 @@ TosaMlirOperatorBuilder::build(TosaSerializationOperator *op) const { return std::vector({mlir_op->getResult(0)}); } +template +mlir::Value +TosaMlirOperatorBuilder::BuildConstShape(mlir::OpBuilder *op_builder, + mlir::Location loc, + const std::vector &values) const { + std::vector 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(loc, type, attr); + block->push_back(mlir_op); + return mlir_op->getResult(0); +} + +template <> +std::vector TosaMlirOperatorBuilder::build( + 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 i64_data; + TosaSerializationHandler::ConvertU8toI64(data, output_type.getRank(), + i64_data); + mlir::Value result = BuildConstShape(op_builder, loc, i64_data); + return std::vector({result}); +} + template std::vector TosaMlirOperatorBuilder::BuildConvOp(TosaSerializationOperator *op) const { @@ -890,6 +963,24 @@ TosaMlirOperatorBuilder::build(TosaSerializationOperator *op) const { return std::vector({mlir_op->getResult(0)}); } +template <> +std::vector TosaMlirOperatorBuilder::build( + TosaSerializationOperator *op) const { + mlir::tosa::shapeType output_type = + shape_type_map->at(op->GetOutputTensorNames()[0]); + + llvm::SmallVector 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( + loc, output_type, input_values); + block->push_back(mlir_op); + return std::vector({mlir_op->getResult(0)}); +} + template <> std::vector TosaMlirOperatorBuilder::build(TosaSerializationOperator *op) const { @@ -926,17 +1017,10 @@ std::vector TosaMlirOperatorBuilder::build( 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(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( - loc, output_type, input_val, new_shape); + loc, output_type, input_val, shape_val); block->push_back(mlir_op); return std::vector({mlir_op->getResult(0)}); } @@ -1081,16 +1165,12 @@ template <> std::vector TosaMlirOperatorBuilder::build(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(op->GetAttribute()); - - mlir::DenseI64ArrayAttr multiples = - BuildDenseI64ArrayAttr(op_builder, attr->multiples()); + Attribute_NONE); // double check attribute type mlir::Operation *mlir_op = op_builder->create( loc, output_type, input_val, multiples); @@ -1395,6 +1475,7 @@ private: TosaMlirRegionBuilder *region_builder; mlir::Block *block; std::unordered_map tensor_type_map; + std::unordered_map shape_type_map; std::unordered_set unranked_tensors; }; @@ -1562,17 +1643,25 @@ mlir::LogicalResult TosaMlirBlockBuilder::BuildAllOpsInBlock( std::queue 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::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 u8_data; + + std::vector data; + auto dense_attr = op.getAttr(llvm::StringRef("value")) + .dyn_cast(); + if (!dense_attr) { + op.emitOpError("Unknown const attribute"); + return nullptr; + } + + for (auto valueIt : dense_attr.getValues()) { + 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::vector{output_name}); + + return tyop; +} + template <> TosaSerializationOperator * TosaSerializationOperatorBuilder::build( @@ -1095,6 +1143,25 @@ TosaSerializationOperatorBuilder::build( return tyop; } +template <> +TosaSerializationOperator * +TosaSerializationOperatorBuilder::build( + mlir::Operation &op) const { + std::vector 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{output_name}); + + return tyop; +} + template <> TosaSerializationOperator * TosaSerializationOperatorBuilder::build( @@ -1122,14 +1189,12 @@ TosaSerializationOperator * TosaSerializationOperatorBuilder::build( 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(op.getAttr("new_shape")); - - TosaReshapeAttribute attribute(shape); TosaSerializationOperator *tyop = new TosaSerializationOperator( - Op_RESHAPE, Attribute_ReshapeAttribute, &attribute, - std::vector{input_name}, + Op_RESHAPE, Attribute_NONE, nullptr, + std::vector{input_name, shape_name}, std::vector{output_name}); return tyop; @@ -1255,15 +1320,13 @@ template <> TosaSerializationOperator * TosaSerializationOperatorBuilder::build( 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(op.getAttr("multiples")); - - TosaTileAttribute attribute(multiples); TosaSerializationOperator *tyop = new TosaSerializationOperator( - Op_TILE, Attribute_TileAttribute, &attribute, - std::vector{input_name}, + Op_TILE, Attribute_NONE, nullptr, + std::vector{input0_name, input1_name}, std::vector{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()) { + auto rank = shape_ty.getRank(); + std::vector shape; + if (rank > 0) { + shape.push_back(rank); + } + ts = new TosaSerializationTensor(name, + /* shape = */ shape, + /* type = */ DType::DType_SHAPE, + /* data = */ std::vector()); + return ts; + } + auto ttype = val.getType().dyn_cast(); 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); } } -- cgit v1.2.1