kevinthesun commented on a change in pull request #4312: [TOPI][Relay][OP]
Dynamic NMS and strided_slice
URL: https://github.com/apache/incubator-tvm/pull/4312#discussion_r345444532
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1947,57 +1948,81 @@ Array<Array<Layout> > StridedSliceInferCorrectLayout(
// NOTE: Discard "const" qualifier here.
auto *params =
const_cast<StridedSliceAttrs*>(attrs.as<StridedSliceAttrs>());
+ CHECK(params != nullptr);
+ Array<Integer> begin, end, strides;
+ const ConstantNode *cbegin, *cend, *cstrides;
+ if ((cbegin = params->begin.as<ConstantNode>()) &&
+ (cend = params->end.as<ConstantNode>()) &&
+ (cstrides = params->strides.as<ConstantNode>())) {
+
+ int32_t* strides_val = reinterpret_cast<int32_t*>(cstrides->data->data);
+ for (size_t i = 0; i < cstrides->data.Shape().front(); ++i){
+ strides.push_back(strides_val[i]);
+ }
+ int32_t* begin_val = reinterpret_cast<int32_t*>(cbegin->data->data);
+ for (size_t i = 0; i < cbegin->data.Shape().front(); ++i){
+ begin.push_back(begin_val[i]);
+ }
+ int32_t* end_val = reinterpret_cast<int32_t*>(cend->data->data);
+ for (size_t i = 0; i < cend->data.Shape().front(); ++i){
+ end.push_back(end_val[i]);
+ }
+ }
Array<Integer> new_begin, new_end;
- for (size_t i = 0; i < params->begin.size(); i++) {
+ for (size_t i = 0; i < begin.size(); i++) {
const LayoutAxis& axis = layout[i];
if (!axis.IsPrimal()) {
// original layout that contains splitted axes is not supported
return {{Layout::Undef()}, {Layout::Undef()}};
}
auto factor = new_layout.FactorOf(axis);
if (factor == -1) {
- new_begin.push_back(params->begin[i]);
- new_end.push_back(params->end[i]);
+ new_begin.push_back(begin[i]);
+ new_end.push_back(end[i]);
} else {
- if (params->strides.defined() && i < params->strides.size()) {
- auto stride = params->strides[i];
+ if (strides.defined() && i < strides.size()) {
+ auto stride = strides[i];
// arbitrary stride is not supported
if (stride.defined() && stride->value != 1) {
return {{Layout::Undef()}, {Layout::Undef()}};
}
}
- int64_t begin = params->begin[i].defined() ? params->begin[i]->value :
0;
- int64_t end = params->end[i].defined() ? params->end[i]->value :
+ int64_t bg = begin[i].defined() ? begin[i]->value : 0;
+ int64_t ed = end[i].defined() ? end[i]->value :
shape[i].as<IntImm>()->value;
- if (begin % factor || end % factor) {
+ if (bg % factor || ed % factor) {
// transform to original layout
return {{Layout::Undef()}, {Layout::Undef()}};
}
- new_begin.push_back(tvm::Integer(begin / factor));
- new_end.push_back(tvm::Integer(end / factor));
+ new_begin.push_back(tvm::Integer(bg / factor));
+ new_end.push_back(tvm::Integer(ed / factor));
}
}
layout = new_layout;
- params->begin = new_begin;
- params->end = new_end;
}
return {{layout}, {layout}};
}
-
-// Positional relay function to create StridedSlice operator used by frontend
FFI.
-Expr MakeStridedSlice(Expr data,
- Array<Integer> begin,
- Array<Integer> end,
- Array<Integer> strides) {
- auto attrs = make_node<StridedSliceAttrs>();
- attrs->begin = std::move(begin);
- attrs->end = std::move(end);
- attrs->strides = std::move(strides);
- static const Op& op = Op::Get("strided_slice");
- return CallNode::make(op, {data}, Attrs(attrs), {});
+inline Tensor DynamicStridedSlice(const tvm::Tensor& input,
Review comment:
Add a TODO here to merge this with existing topi StridedSlice?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
