MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / apply_on_physical_tensor

Function apply_on_physical_tensor

imperative/src/impl/ops/subtensor.cpp:176–218  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

174}
175
176SmallVector<TensorPtr> apply_on_physical_tensor(
177 const OpDef& def, const SmallVector<TensorPtr>& inputs,
178 SmallVector<LogicalTensorDesc>& output_descs, const bool& validated) {
179 CompNode cn = inputs[0]->comp_node();
180 auto&& layout = inputs[0]->layout();
181 auto&& op = static_cast<const Subtensor&>(def);
182
183 if (inputs.size() > 1) {
184 return proxy_graph_detail::apply_on_physical_tensor(
185 def, inputs, output_descs, validated);
186 }
187 auto&& src = inputs[0];
188 auto slice_items = op.slice_items;
189 auto items = op.items;
190 TensorLayout res_layout = deduce_layout(layout, items, slice_items);
191 if (res_layout.is_empty()) {
192 return {Tensor::make(res_layout, cn)};
193 }
194 size_t offset = 0;
195 size_t dtype_size = layout.dtype.size();
196 TensorPtr tensor = src;
197 for (int i = items.size() - 1; i >= 0; i--) {
198 auto&& [axis, b_flag, e_flag, s_flag, idx_flag] = items[i];
199 auto&& [b_val, e_val, s_val, ax_val] = slice_items[i];
200 int start = b_val;
201 if (idx_flag) {
202 ax_val = ax_val < 0 ? layout.shape[axis] + ax_val : ax_val;
203 offset += ax_val * layout.stride[axis] * dtype_size;
204 } else {
205 int shape_axis = src->layout().shape[axis];
206 if (s_val < 0) {
207 start = b_val == INT_MIN ? shape_axis - 1 : b_val;
208 start = mod_size(start, shape_axis);
209 }
210 start = start == INT_MIN ? 0 : start;
211 start = start < 0 ? start + shape_axis : start;
212 offset += start * layout.stride[axis] * dtype_size;
213 }
214 }
215
216 // memory forward
217 return {Tensor::make(src->blob(), src->offset() + offset, res_layout)};
218}
219
220SmallVector<VarNode::LayoutConstraintCallback> get_input_layout_constraint(
221 const OpDef& def, const SmallVector<TensorPtr>& inputs) {

Callers

nothing calls this directly

Calls 8

deduce_layoutFunction · 0.85
blobMethod · 0.80
makeFunction · 0.70
comp_nodeMethod · 0.45
layoutMethod · 0.45
sizeMethod · 0.45
is_emptyMethod · 0.45
offsetMethod · 0.45

Tested by

no test coverage detected