| 31 | } |
| 32 | |
| 33 | void apply_on_device_tensornd( |
| 34 | const OpDef& def, const SmallVector<DeviceTensorND>& inputs, |
| 35 | SmallVector<DeviceTensorND>* outputs) { |
| 36 | auto&& op_def = def.cast_final_safe<GetVarShape>(); |
| 37 | |
| 38 | TensorShape shp; |
| 39 | if (inputs.size() == 1) { |
| 40 | shp = inputs[0].layout(); |
| 41 | } else { |
| 42 | TensorShapeArray src(inputs.size()); |
| 43 | for (size_t i = 0; i < inputs.size(); ++i) { |
| 44 | src[i] = inputs[i].layout(); |
| 45 | } |
| 46 | megdnn::Elemwise::deduce_shape(src, shp); |
| 47 | } |
| 48 | |
| 49 | mgb_assert(shp.ndim != 0, "input shape invalid"); |
| 50 | mgb_assert( |
| 51 | (*outputs)[0].comp_node() == CompNode::default_cpu(), |
| 52 | "GetVarShape's apply_on_device_tensornd should receive default_cpu " |
| 53 | "outputs."); |
| 54 | |
| 55 | HostTensorND hv; |
| 56 | if (op_def.axis == opr::GetVarShape::Param::INVALID_AXIS) { |
| 57 | hv = HostTensorND(CompNode::default_cpu(), {shp.ndim}, dtype::Int32()); |
| 58 | auto* ptr = hv.ptr<dt_int32>(); |
| 59 | for (size_t i = 0; i < shp.ndim; ++i) { |
| 60 | ptr[i] = shp.shape[i]; |
| 61 | } |
| 62 | } else { |
| 63 | int32_t axis = op_def.axis; |
| 64 | if (axis < 0) { |
| 65 | axis += shp.ndim; |
| 66 | } |
| 67 | mgb_assert(axis >= 0 && axis < (int32_t)shp.ndim); |
| 68 | hv = HostTensorND(CompNode::default_cpu(), {1}, dtype::Int32()); |
| 69 | auto* ptr = hv.ptr<dt_int32>(); |
| 70 | ptr[0] = shp.shape[axis]; |
| 71 | } |
| 72 | (*outputs)[0] = DeviceTensorND::make_proxy(hv); |
| 73 | } |
| 74 | |
| 75 | HostTensorND get_var_shape_host_tensor( |
| 76 | const OpDef& def, const SmallVector<TensorPtr>& inputs) { |
no test coverage detected