| 73 | } |
| 74 | |
| 75 | DispatchMode decide_dispatch_mode( |
| 76 | const OpDef& def, const SmallVector<LogicalTensorDesc>& inputs) { |
| 77 | bool host_computable = true; |
| 78 | constexpr int size_threshhold = TensorShape::MAX_NDIM; |
| 79 | for (auto&& inp : inputs) { |
| 80 | if (inp.value.empty() || inp.value.layout().ndim == 0 || |
| 81 | inp.value.layout().total_nr_elems() > size_threshhold) { |
| 82 | host_computable = false; |
| 83 | break; |
| 84 | } |
| 85 | } |
| 86 | return host_computable ? DEFAULT_CPU : KERNEL; |
| 87 | } |
| 88 | |
| 89 | void apply_on_device_tensornd( |
| 90 | const OpDef& def, const SmallVector<DeviceTensorND>& inputs, |
no test coverage detected