colin
/
TIM-VX
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add data layout for kernel to support TVM conv2d (#40) 

Signed-off-by: Zongwu Yang <zongwu.yang@verisilicon.com>
This commit is contained in:
Zongwu.Yang 2021-05-14 14:00:02 +08:00 committed by GitHub
parent 096d99e068
commit b38cad9f1d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 89 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
include/tim/vx/ops/conv2d.h
View File

@ -38,13 +38,17 @@ class Conv2d : public Operation {
const std::array<uint32_t, 2>& ksize,
const std::array<uint32_t, 2>& stride,
const std::array<uint32_t, 2>& dilation, int32_t multiplier = 0,
DataLayout layout = DataLayout::WHCN);
DataLayout input_layout = DataLayout::WHCN,
DataLayout kernel_layout = DataLayout::WHIcOc);
Conv2d(Graph* graph, int32_t weights, PadType padding,
const std::array<uint32_t, 2>& ksize,
const std::array<uint32_t, 2>& stride,
const std::array<uint32_t, 2>& dilation,
const std::array<uint32_t, 4>& pad, int32_t multiplier = 0,
DataLayout layout = DataLayout::WHCN);
DataLayout input_layout = DataLayout::WHCN,
DataLayout kernel_layout = DataLayout::WHIcOc);
DataLayout KernelDataLayout() { return kernel_layout_; }
protected:
const uint32_t weights_;
@ -54,6 +58,7 @@ class Conv2d : public Operation {
const std::array<uint32_t, 2> dilation_;
const std::array<uint32_t, 4> pad_;
const int32_t multiplier_;
const DataLayout kernel_layout_;
};
} // namespace ops

9
include/tim/vx/types.h
View File

@ -74,7 +74,14 @@ enum class ActivationType {
enum class ResizeType { NEAREST_NEIGHBOR, BILINEAR, AREA };
enum class DataLayout { WHCN, CWHN, ANY };
enum class DataLayout {
WHCN,
CWHN,
ANY,
IcWHOc /*TF*/,
OcIcWH /*TVM*/,
WHIcOc /*TIM-VX default*/
};
} // namespace vx
} // namespace tim

16
src/tim/transform/ops/conv2d_layout_inference.h
View File

@ -61,8 +61,16 @@ class Conv2dLayoutInfer : public OpLayoutInfer {
} else {
// For input/weight
if (!required_pv->IsAligned()) {
auto src_conv2d = std::static_pointer_cast<vx::ops::Conv2d>(op_);
// Support TVM Kernel Layout
if (src_conv2d->KernelDataLayout() == vx::DataLayout::OcIcWH) {
trans_pv = std::make_shared<PermuteVector<4>>(kOcIcWH2WHIcOc);
infer_tensor = PermuteConstTensor(
in, trans_pv);
} else {
infer_tensor = PermuteConstTensor(in, required_pv);
trans_pv = required_pv;
}
} else {
infer_tensor = context_->infer_graph_->CreateTensor(
in->GetSpec(), in->GetDataRef());
@ -108,11 +116,11 @@ class Conv2dLayoutInfer : public OpLayoutInfer {
int32_t out_channels = op_->impl()->node()->nn_param.conv2d.weights;
auto conv2d = context_->infer_graph_->CreateOperation<vx::ops::Conv2d>(
out_channels, pad_type, ksize, stride, dilation, pad, multiplier,
vx::DataLayout::WHCN);
vx::DataLayout::WHCN, vx::DataLayout::WHIcOc);
auto otensor_infer = CreateOutputsTensor(required_pv);
(*conv2d).BindInputs({context_->GetMapedTensor(input_tensors[0]),
context_->GetMapedTensor(input_tensors[1]),
context_->GetMapedTensor(input_tensors[2])});
for (const auto& i_src : input_tensors) {
(*conv2d).BindInput(context_->GetMapedTensor(i_src));
}
(*conv2d).BindOutput(otensor_infer[0]);
context_->SetPermuteVector(op_->impl()->OutputsTensor()[0], required_pv);
// Add out tensor of src_graph into next_tensor

55
src/tim/transform/ops/op_layout_inference.cc
View File

@ -154,16 +154,36 @@ OpLayoutInfer::AlignPermuteVectorForMutilInputs() {
auto src_inputs = op_->impl()->InputsTensor();
// Suppose the inputs have same dimension rank
// TODO(yzw): should choose a optimal required_pv
auto required_pv = context_->GetPermuteVector(src_inputs[0]);
std::shared_ptr<IPermuteVector> required_pv = nullptr;
for (const auto& in : src_inputs) {
if (!in->IsConstTensor()) {
required_pv = context_->GetPermuteVector(in);
break;
}
}
if (!required_pv) {
// all inputs are constant tensors
for (const auto& i_src : src_inputs) {
context_->UpdateTensorMap(
i_src, context_->infer_graph_->CreateTensor(i_src->GetSpec(),
i_src->GetDataRef()));
context_->SetPermuteVector(i_src, MakeShared(i_src->GetShape().size()));
}
} else {
for (const auto& i_src : src_inputs) {
std::shared_ptr<vx::Tensor> perm_out;
auto pv = context_->GetPermuteVector(i_src);
auto final_pv = pv->Reverse()->Add(required_pv);
if (!final_pv->IsAligned()) {
if (i_src->IsConstTensor()) {
perm_out = PermuteConstTensor(i_src, final_pv);
required_pv->IsAligned()
? perm_out = context_->infer_graph_->CreateTensor(i_src->GetSpec(),
i_src->GetDataRef())
: perm_out = PermuteConstTensor(i_src, required_pv);
} else {
perm_out = InsertPermute(context_->GetMapedTensor(i_src), final_pv);
auto final_pv =
context_->GetPermuteVector(i_src)->Reverse()->Add(required_pv);
final_pv->IsAligned() ? perm_out = context_->GetMapedTensor(i_src)
: perm_out = InsertPermute(
context_->GetMapedTensor(i_src), final_pv);
}
context_->UpdateTensorMap(i_src, perm_out);
context_->SetPermuteVector(i_src, required_pv);
@ -175,18 +195,22 @@ OpLayoutInfer::AlignPermuteVectorForMutilInputs() {
void OpLayoutInfer::ReverseInputsPermuteVector() {
for (const auto& i_src : op_->impl()->InputsTensor()) {
std::shared_ptr<vx::Tensor> perm_out;
auto input_pv = context_->GetPermuteVector(i_src);
if (!input_pv->IsAligned()) {
std::shared_ptr<IPermuteVector> input_pv;
if (i_src->IsConstTensor()) {
perm_out = PermuteConstTensor(i_src, input_pv);
perm_out = context_->infer_graph_->CreateTensor(i_src->GetSpec(),
i_src->GetDataRef());
input_pv = MakeShared(i_src->GetShape().size());
} else {
perm_out = context_->GetMapedTensor(i_src);
input_pv = context_->GetPermuteVector(i_src);
if (!input_pv->IsAligned()) {
perm_out =
InsertPermute(context_->GetMapedTensor(i_src), input_pv->Reverse());
InsertPermute(perm_out, input_pv->Reverse());
}
}
context_->UpdateTensorMap(i_src, perm_out);
context_->SetPermuteVector(i_src, MakeShared(input_pv->Rank()));
}
}
}
bool OpLayoutInfer::TransposeConstTensorData(
@ -202,12 +226,15 @@ bool OpLayoutInfer::TransposeConstTensorData(
return false;
}
std::vector<uint32_t> perm = KOcHWIc2OcIcHW;
vx::ShapeType reverse_shape;
for (int32_t i = input->GetShape().size() - 1; i >= 0; i--) {
reverse_shape.push_back(input->GetShape()[i]);
}
std::vector<uint32_t> perm = KOcHWIc2OcIcHW;
std::vector<uint32_t>tmp_vec = kOcIcWH2WHIcOc;
if (pv->AsStdVec() == tmp_vec) {
perm = kHWIcOc2OcIcHW;
}
vsi_nn_Transpose(out_data.data(), (uint8_t*)(input->GetDataRef()),
(uint32_t*)(reverse_shape.data()),
static_cast<uint32_t>(input->GetShape().size()),
@ -240,7 +267,7 @@ std::vector<uint32_t> OpLayoutInfer::MapPadding(const std::vector<uint32_t>& per
assert(perm.size() == padding.size());
std::vector<uint32_t> r(padding.size());
for (int i = 0; i < padding.size(); ++i) {
for (uint32_t i = 0; i < padding.size(); ++i) {
r[i] = padding[perm[i]];
}

5
src/tim/transform/ops/op_layout_inference.h
View File

@ -34,7 +34,12 @@ namespace tim {
namespace transform {
constexpr std::initializer_list<uint32_t> kCWHN2WHCN = {1, 2, 0, 3};
constexpr std::initializer_list<uint32_t> KOcHWIc2OcIcHW = {0, 3, 1, 2};
constexpr std::initializer_list<uint32_t> kIcWHOc2WHIcOc = {1, 2, 0, 3};
constexpr std::initializer_list<uint32_t> kHWIcOc2OcIcHW = {3, 2, 0, 1};
constexpr std::initializer_list<uint32_t> kOcIcWH2WHIcOc = {2, 3, 1, 0};
class OpLayoutInfer {
public:

16
src/tim/vx/ops/conv2d.cc
View File

@ -34,24 +34,26 @@ namespace ops {
Conv2d::Conv2d(Graph* graph, int32_t weights, PadType padding,
const std::array<uint32_t, 2>& ksize,
const std::array<uint32_t, 2>& stride,
const std::array<uint32_t, 2>& dilation,
int32_t multiplier, DataLayout layout)
: Conv2d(graph, weights, padding, ksize, stride, dilation,
{0, 0, 0, 0}, multiplier, layout) {}
const std::array<uint32_t, 2>& dilation, int32_t multiplier,
DataLayout input_layout, DataLayout kernel_layout)
: Conv2d(graph, weights, padding, ksize, stride, dilation, {0, 0, 0, 0},
multiplier, input_layout, kernel_layout) {}
Conv2d::Conv2d(Graph* graph, int32_t weights, PadType padding,
const std::array<uint32_t, 2>& ksize,
const std::array<uint32_t, 2>& stride,
const std::array<uint32_t, 2>& dilation,
const std::array<uint32_t, 4>& pad, int32_t multiplier, DataLayout layout)
: Operation(graph, VSI_NN_OP_CONV2D, 0, 0, layout),
const std::array<uint32_t, 4>& pad, int32_t multiplier,
DataLayout input_layout, DataLayout kernel_layout)
: Operation(graph, VSI_NN_OP_CONV2D, 0, 0, input_layout),
weights_(weights),
padding_(padding),
ksize_(ksize),
stride_(stride),
dilation_(dilation),
pad_(pad),
multiplier_(multiplier) {
multiplier_(multiplier),
kernel_layout_(kernel_layout) {
this->impl()->node()->nn_param.conv2d.ksize[0] = ksize_[0];
this->impl()->node()->nn_param.conv2d.ksize[1] = ksize_[1];
this->impl()->node()->nn_param.conv2d.stride[0] = stride_[0];