commit
4f4fa5260c
@ -0,0 +1,36 @@
|
||||
/**
|
||||
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include "src/ops/select.h"
|
||||
#include "src/ops/primitive_c.h"
|
||||
#include "src/ops/populate/populate_register.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace lite {
|
||||
OpParameter *PopulateSelectParameter(const mindspore::lite::PrimitiveC *primitive) {
|
||||
OpParameter *select_parameter = reinterpret_cast<OpParameter *>(malloc(sizeof(OpParameter)));
|
||||
if (select_parameter == nullptr) {
|
||||
MS_LOG(ERROR) << "malloc SelectParameter failed.";
|
||||
return nullptr;
|
||||
}
|
||||
memset(select_parameter, 0, sizeof(OpParameter));
|
||||
select_parameter->type_ = primitive->Type();
|
||||
|
||||
return reinterpret_cast<OpParameter *>(select_parameter);
|
||||
}
|
||||
Registry SelectParameterRegistry(schema::PrimitiveType_Select, PopulateSelectParameter);
|
||||
} // namespace lite
|
||||
} // namespace mindspore
|
@ -0,0 +1,106 @@
|
||||
/**
|
||||
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#include "src/ops/select.h"
|
||||
|
||||
#ifndef PRIMITIVE_WRITEABLE
|
||||
#include "src/ops/ops_register.h"
|
||||
#endif
|
||||
#include "src/tensorlist.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace lite {
|
||||
#ifdef PRIMITIVE_WRITEABLE
|
||||
int Select::UnPackAttr(const Primitive &prim, const std::vector<AnfNodePtr> &inputs) {
|
||||
if (this->primitive_ == nullptr) {
|
||||
this->primitive_ = new (std::nothrow) schema::PrimitiveT;
|
||||
if (this->primitive_ == nullptr) {
|
||||
MS_LOG(ERROR) << "new primitiveT failed";
|
||||
return RET_ERROR;
|
||||
}
|
||||
this->primitive_->value.type = schema::PrimitiveType_Select;
|
||||
}
|
||||
if (this->primitive_->value.type != schema::PrimitiveType_Select) {
|
||||
MS_LOG(ERROR) << "Primitive type is error :" << this->primitive_->value.type;
|
||||
return RET_ERROR;
|
||||
}
|
||||
if (this->primitive_->value.value == nullptr) {
|
||||
auto attr = new (std::nothrow) schema::SelectT();
|
||||
if (attr == nullptr) {
|
||||
MS_LOG(ERROR) << "new primitiveT value failed";
|
||||
return RET_ERROR;
|
||||
}
|
||||
|
||||
this->primitive_->value.value = attr;
|
||||
if (this->primitive_->value.value == nullptr) {
|
||||
MS_LOG(ERROR) << "primitive value is nullptr";
|
||||
return RET_ERROR;
|
||||
}
|
||||
}
|
||||
return RET_OK;
|
||||
}
|
||||
#else
|
||||
int Select::UnPackToFlatBuilder(const schema::Primitive *primitive, flatbuffers::FlatBufferBuilder *fbb) {
|
||||
MS_ASSERT(nullptr != primitive);
|
||||
MS_ASSERT(nullptr != fbb);
|
||||
auto attr = primitive->value_as_Select();
|
||||
if (attr == nullptr) {
|
||||
MS_LOG(ERROR) << "value_as_Select return nullptr";
|
||||
return RET_ERROR;
|
||||
}
|
||||
auto val_offset = schema::CreateSelect(*fbb);
|
||||
auto prim_offset = schema::CreatePrimitive(*fbb, schema::PrimitiveType_Select, val_offset.o);
|
||||
fbb->Finish(prim_offset);
|
||||
return RET_OK;
|
||||
}
|
||||
|
||||
PrimitiveC *SelectCreator(const schema::Primitive *primitive) { return PrimitiveC::NewPrimitiveC<Select>(primitive); }
|
||||
Registry SelectRegistry(schema::PrimitiveType_Select, SelectCreator);
|
||||
#endif
|
||||
|
||||
int Select::InferShape(std::vector<Tensor *> inputs_, std::vector<Tensor *> outputs_) {
|
||||
MS_ASSERT(inputs_.size() == 2 * outputs_.size() + 1);
|
||||
if (!infer_flag()) {
|
||||
return RET_INFER_INVALID;
|
||||
}
|
||||
for (size_t i = 0; i < outputs_.size(); i++) {
|
||||
auto *input = inputs_[i + 1];
|
||||
auto *output = outputs_[i];
|
||||
if (input == nullptr) {
|
||||
MS_LOG(ERROR) << "input tensor is nullptr";
|
||||
return RET_ERROR;
|
||||
}
|
||||
if (output == nullptr) {
|
||||
MS_LOG(ERROR) << "output tensor is nullptr";
|
||||
return RET_ERROR;
|
||||
}
|
||||
output->set_data_type(input->data_type());
|
||||
output->set_shape(input->shape());
|
||||
output->set_format(input->format());
|
||||
auto data_type = input->data_type();
|
||||
if (data_type != kObjectTypeTensorType) {
|
||||
continue;
|
||||
} else {
|
||||
auto input_tensorlist = reinterpret_cast<TensorList *>(input);
|
||||
auto output_tensorlist = reinterpret_cast<TensorList *>(output);
|
||||
output_tensorlist->set_element_shape(input_tensorlist->element_shape());
|
||||
output_tensorlist->set_max_elements_num(input_tensorlist->max_elements_num());
|
||||
output_tensorlist->set_tensors_data_type(input_tensorlist->tensors_data_type());
|
||||
}
|
||||
}
|
||||
return RET_OK;
|
||||
}
|
||||
} // namespace lite
|
||||
} // namespace mindspore
|
@ -0,0 +1,66 @@
|
||||
/**
|
||||
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#include "src/runtime/kernel/arm/base/select.h"
|
||||
#include "src/kernel_registry.h"
|
||||
#include "include/errorcode.h"
|
||||
#include "src/tensorlist.h"
|
||||
|
||||
using mindspore::lite::KernelRegistrar;
|
||||
using mindspore::lite::RET_ERROR;
|
||||
using mindspore::lite::RET_OK;
|
||||
using mindspore::schema::PrimitiveType_Select;
|
||||
|
||||
namespace mindspore::kernel {
|
||||
int SelectCPUKernel::Init() { return RET_OK; }
|
||||
|
||||
int SelectCPUKernel::ReSize() { return RET_OK; }
|
||||
|
||||
// inputs: bool*1 true-data*n false-data*n
|
||||
// output: data*n
|
||||
int SelectCPUKernel::Run() {
|
||||
MS_ASSERT(in_tensors_.size() >= 3);
|
||||
MS_ASSERT(in_tensors_.size() == out_tensors_.size() * 2 + 1);
|
||||
auto bool_tensor = in_tensors_.front();
|
||||
MS_ASSERT(bool_tensor != nullptr);
|
||||
MS_ASSERT(bool_tensor->data_type() == kNumberTypeBool);
|
||||
MS_ASSERT(bool_tensor->Size() == 1);
|
||||
MS_ASSERT(bool_tensor->Size() == 1);
|
||||
auto condition = static_cast<bool *>(bool_tensor->data_c());
|
||||
if (condition == nullptr) {
|
||||
MS_LOG(ERROR) << "data of bool tensor is nullptr";
|
||||
return lite::RET_NULL_PTR;
|
||||
}
|
||||
if (*condition) {
|
||||
auto ret = MoveData(this->out_tensors_.begin(), this->out_tensors_.end(), this->in_tensors_.begin() + 1,
|
||||
this->in_tensors_.begin() + 1 + this->out_tensors_.size());
|
||||
if (ret != RET_OK) {
|
||||
MS_LOG(ERROR) << "carry data error : " << ret;
|
||||
return ret;
|
||||
}
|
||||
} else {
|
||||
auto ret = MoveData(this->out_tensors_.begin(), this->out_tensors_.end(),
|
||||
this->in_tensors_.begin() + 1 + this->out_tensors_.size(),
|
||||
this->in_tensors_.begin() + 1 + 2 * this->out_tensors_.size());
|
||||
if (ret != RET_OK) {
|
||||
MS_LOG(ERROR) << "carry data error : " << ret;
|
||||
return ret;
|
||||
}
|
||||
}
|
||||
return RET_OK;
|
||||
}
|
||||
|
||||
REG_KERNEL(kCPU, kNumberTypeBool, PrimitiveType_Select, LiteKernelCreator<SelectCPUKernel>)
|
||||
} // namespace mindspore::kernel
|
@ -0,0 +1,38 @@
|
||||
/**
|
||||
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#ifndef MINDSPORE_LITE_SRC_RUNTIME_KERNEL_ARM_BASE_SELECT_H_
|
||||
#define MINDSPORE_LITE_SRC_RUNTIME_KERNEL_ARM_BASE_SELECT_H_
|
||||
|
||||
#include <vector>
|
||||
#include "src/runtime/kernel/arm/base/carry_data.h"
|
||||
#include "src/lite_kernel.h"
|
||||
#include "src/tensorlist.h"
|
||||
|
||||
namespace mindspore::kernel {
|
||||
class SelectCPUKernel : public CarryDataKernel {
|
||||
public:
|
||||
SelectCPUKernel(OpParameter *parameter, const std::vector<lite::Tensor *> &inputs,
|
||||
const std::vector<lite::Tensor *> &outputs, const lite::InnerContext *ctx,
|
||||
const mindspore::lite::PrimitiveC *primitive)
|
||||
: CarryDataKernel(parameter, inputs, outputs, ctx, primitive) {}
|
||||
~SelectCPUKernel() override = default;
|
||||
int Init() override;
|
||||
int ReSize() override;
|
||||
int Run() override;
|
||||
};
|
||||
} // namespace mindspore::kernel
|
||||
|
||||
#endif // MINDSPORE_LITE_SRC_RUNTIME_KERNEL_ARM_BASE_SELECT_H_
|
Loading…
Reference in new issue