| /external/pytorch/aten/src/ATen/core/dispatch/ |
| D | OperatorEntry.cpp | 106 std::optional<DispatchKey> dispatch_key, in registerKernel() argument
127 " dispatch key: ", toString(local_cpp_signature->dispatch_key), "\n", in registerKernel()
130 " dispatch key: ", toString(dispatch_key), "\n", in registerKernel()
134 local_cpp_signature = CppSignatureWithDebug { *cpp_signature, debug, dispatch_key }; in registerKernel()
145 …auto& k = dispatch_key.has_value() ? kernels_[*dispatch_key] : kernels_[DispatchKey::CompositeImpl… in registerKernel()
154 if (dispatch_key != DispatchKey::Meta) { in registerKernel()
159 " dispatch key: ", toString(dispatch_key), "\n", in registerKernel()
176 if (dispatch_key.has_value()) { in registerKernel()
177 updateDispatchTable_(dispatcher, *dispatch_key); in registerKernel()
186 std::optional<DispatchKey> dispatch_key, in registerKernel() argument
[all …]
|
| D | OperatorEntry.h | 132 std::optional<DispatchKey> dispatch_key,
142 std::optional<DispatchKey> dispatch_key,
149 DispatchKey dispatch_key
285 std::optional<DispatchKey> dispatch_key; member
297 …tion& computeDispatchTableEntry(const c10::Dispatcher& dispatcher, DispatchKey dispatch_key) const;
299 const c10::Dispatcher& dispatcher, DispatchKey dispatch_key
303 void updateDispatchTableEntry_(const c10::Dispatcher& dispatcher, DispatchKey dispatch_key);
305 void updateDispatchTable_(const c10::Dispatcher& dispatcher, DispatchKey dispatch_key);
309 const AnnotatedKernel* getKernelForDispatchKey(DispatchKey dispatch_key) const;
|
| /external/pytorch/torch/csrc/utils/ |
| D | tensor_new.h | 34 c10::DispatchKey dispatch_key,
39 c10::DispatchKey dispatch_key,
49 c10::DispatchKey dispatch_key,
53 c10::DispatchKey dispatch_key,
59 c10::DispatchKey dispatch_key,
63 c10::DispatchKey dispatch_key,
67 c10::DispatchKey dispatch_key,
71 c10::DispatchKey dispatch_key,
75 c10::DispatchKey dispatch_key,
80 c10::DispatchKey dispatch_key,
[all …]
|
| D | tensor_new.cpp | 538 c10::DispatchKey dispatch_key, in check_base_legacy_new() argument
555 expected_key_set.has(dispatch_key), in check_base_legacy_new()
559 dispatch_key); in check_base_legacy_new()
570 expected_key_set.has(dispatch_key), in check_base_legacy_new()
574 dispatch_key); in check_base_legacy_new()
582 c10::DispatchKey dispatch_key, in check_legacy_ctor_device() argument
586 dispatchKeyToDeviceType(dispatch_key) == device.value().type(), in check_legacy_ctor_device()
588 dispatchKeyToDeviceType(dispatch_key), in check_legacy_ctor_device()
602 c10::DispatchKey dispatch_key, in legacy_sparse_tensor_generic_ctor_new() argument
607 auto options = dispatchKeyToTensorOptions(dispatch_key); in legacy_sparse_tensor_generic_ctor_new()
[all …]
|
| D | nested.cpp | 18 c10::DispatchKey dispatch_key) { in typeIdWithDefault() argument
19 auto options = dispatchKeyToTensorOptions(dispatch_key); in typeIdWithDefault()
27 c10::DispatchKey dispatch_key, in nested_tensor_ctor() argument
40 auto tensor_options = typeIdWithDefault(r, 2, dispatch_key); in nested_tensor_ctor()
71 new_list[i] = tensor_ctor(dispatch_key, scalar_type, elem_r); in nested_tensor_ctor()
|
| D | python_dispatch.cpp | 117 c10::DispatchKey dispatch_key, in PythonKernelHolder() argument
121 dispatch_key_(dispatch_key), in PythonKernelHolder()
649 [](c10::DispatchKey dispatch_key, bool desired_state) { in initDispatchBindings() argument
650 c10::impl::tls_set_dispatch_key_excluded(dispatch_key, desired_state); in initDispatchBindings()
654 [](c10::DispatchKey dispatch_key) { in initDispatchBindings() argument
655 return c10::impl::tls_is_dispatch_key_excluded(dispatch_key); in initDispatchBindings()
659 [](c10::DispatchKey dispatch_key, bool desired_state) { in initDispatchBindings() argument
660 c10::impl::tls_set_dispatch_key_included(dispatch_key, desired_state); in initDispatchBindings()
664 [](c10::DispatchKey dispatch_key) { in initDispatchBindings() argument
665 return c10::impl::tls_is_dispatch_key_included(dispatch_key); in initDispatchBindings()
[all …]
|
| /external/pytorch/torch/ |
| D | library.py | 75 def __init__(self, ns, kind, dispatch_key=""): argument
88 kind, ns, dispatch_key, filename, lineno
95 self.dispatch_key = dispatch_key
208 def _impl_with_aoti_compile(self, op_name, dispatch_key=""): argument
220 if dispatch_key == "":
221 dispatch_key = self.dispatch_key
222 assert torch.DispatchKeySet(dispatch_key).has(torch._C.DispatchKey.Dense)
237 key = self.ns + "/" + name.split("::")[-1] + "/" + dispatch_key
244 name.split("::")[-1], dispatch_key, self.ns
250 impl_fn(self.ns, name.split("::")[-1], dispatch_key)
[all …]
|
| D | _ops.py | 264 for dispatch_key in _HIGHER_ORDER_OP_DEFAULT_FALLTHROUGH_DISPATCH_KEYS:
265 self.fallthrough(dispatch_key)
284 def fallthrough(self, dispatch_key): argument
285 self.non_fallthrough_keys = self.non_fallthrough_keys.remove(dispatch_key)
289 def dispatch(self, /, dispatch_key, *args, **kwargs): argument
292 if dispatch_key in self._dispatch_cache:
293 kernel = self._dispatch_cache[dispatch_key]
297 if dispatch_key == DispatchKey.FuncTorchDynamicLayerFrontMode:
300 if dispatch_key == DispatchKey.Python:
370 … functionality_key = torch._C._to_functionality_key(dispatch_key) # type: ignore[attr-defined]
[all …]
|
| /external/pytorch/torchgen/dest/ |
| D | register_dispatch_key.py | 56 if backend_index.dispatch_key in (DispatchKey.CPU, DispatchKey.Meta):
58 elif backend_index.dispatch_key == DispatchKey.CUDA:
63 elif backend_index.dispatch_key == DispatchKey.MPS:
65 elif backend_index.dispatch_key == DispatchKey.XPU:
88 if backend_index.dispatch_key in (
95 dispatch = str(backend_index.dispatch_key).lower()
98 elif backend_index.dispatch_key in (
111 if backend_index.dispatch_key == DispatchKey.Meta:
152 if backend_index.dispatch_key == DispatchKey.CompositeExplicitAutogradNonFunctional:
365 if self.backend_index.dispatch_key == DispatchKey.Meta:
[all …]
|
| /external/pytorch/torchgen/ |
| D | gen_aoti_c_shim.py | 317 dispatch_key: DispatchKey,
321 if backend_indices[dispatch_key].has_kernel(func) or (
324 and backend_indices[dispatch_key].has_kernel(
328 backend_index = backend_indices[dispatch_key]
348 dispatch_key: DispatchKey,
352 func, func_group_mapping, dispatch_key, backend_indices
372 dispatch_key: DispatchKey,
377 func, func_group_mapping, dispatch_key, backend_indices
383 device = dispatch_key.lower()
406 dispatch_key: DispatchKey
[all …]
|
| D | gen_backend_stubs.py | 126 dispatch_key: DispatchKey,
147 dispatch_key=dispatch_key,
421 dispatch_key: DispatchKey,
438 backend_index = backend_indices[dispatch_key]
466 dispatch_key=dispatch_key,
480 dispatch_key=dispatch_key,
493 "DispatchKey": dispatch_key,
494 "dispatch_namespace": dispatch_key.lower(),
506 "dispatch_namespace": dispatch_key.lower(),
595 for dispatch_key in (
[all …]
|
| D | gen.py | 170 dispatch_key=DispatchKey.Undefined,
184 dispatch_key=k,
327 return [backend.dispatch_key for backend in backends] + [
343 return backend_index.dispatch_key
363 dispatch_key = get_static_dispatch_backend(f, index)
364 if dispatch_key is not None:
374 for dispatch_key in static_dispatch_keys(backends)
487 and b.dispatch_key in STRUCTURED_DISPATCH_KEYS
956 dispatch_key = "c10::computeDispatchKey(dtype, layout, device)"
1442 for dispatch_key, backend_idx in backend_indices.items():
[all …]
|
| /external/pytorch/aten/src/ATen/core/op_registration/ |
| D | op_registration.h | 72 Options&& kernel(DispatchKey dispatch_key) && { in kernel() argument
73 …return std::move(*this).kernel(dispatch_key, KernelFunction::makeFromBoxedFunction<kernel_func>(),… in kernel()
161 …guts::is_functor<KernelFunctor>::value, Options&&> kernel(DispatchKey dispatch_key, ConstructorPar… in kernel() argument
166 dispatch_key, in kernel()
241 …able_if_t<guts::is_function_type<FuncType>::value, Options&&> kernel(DispatchKey dispatch_key) && { in kernel() argument
246 dispatch_key, in kernel()
285 …uts::is_function_type<FuncType>::value, Options&&> kernel(DispatchKey dispatch_key, FuncType* kern… in kernel() argument
290 dispatch_key, in kernel()
334 Options&&> kernel(DispatchKey dispatch_key, Lambda&& functor) && { in kernel() argument
346 dispatch_key, in kernel()
[all …]
|
| D | op_registration.cpp | 97 if (kernel.dispatch_key.has_value()) { in checkNoDuplicateKernels_()
99 0 == dispatch_keys.count(*kernel.dispatch_key), in checkNoDuplicateKernels_()
101 *kernel.dispatch_key, in checkNoDuplicateKernels_()
104 dispatch_keys.insert(*kernel.dispatch_key); in checkNoDuplicateKernels_()
133 kernel.dispatch_key, in registerOp_()
|
| /external/pytorch/aten/src/ATen/core/ |
| D | library.cpp | 185 auto dispatch_key = f.dispatch_key_.has_value() ? f.dispatch_key_ : dispatch_key_; in _def() local
189 dispatch_key, in _def()
238 auto dispatch_key = f.dispatch_key_.has_value() ? f.dispatch_key_ : dispatch_key_; in _impl() local
244 dispatch_key, in _impl()
253 c10::Dispatcher::singleton().waitForImpl(name, dispatch_key); in _impl()
270 auto dispatch_key = f.dispatch_key_.has_value() ? f.dispatch_key_ : dispatch_key_; in _fallback() local
271 TORCH_INTERNAL_ASSERT(dispatch_key.has_value(), ERROR_CONTEXT); in _fallback()
276 " TORCH_LIBRARY_IMPL(_, ", *dispatch_key, ", m) { m.fallback(...); }\n\n", in _fallback()
280 for (auto k : c10::getRuntimeDispatchKeySet(*dispatch_key)) { in _fallback()
|
| /external/pytorch/aten/src/ATen/core/boxing/impl/ |
| D | kernel_stackbased_test.cpp | 56 void expectCallsIncrement(DispatchKey dispatch_key) { in expectCallsIncrement() argument
57 expectCallsIncrement(c10::DispatchKeySet(dispatch_key)); in expectCallsIncrement()
60 void expectCallsIncrementUnboxed(DispatchKey dispatch_key) { in expectCallsIncrementUnboxed() argument
66 int64_t result = callOpUnboxed<int64_t, at::Tensor, int64_t>(*op, dummyTensor(dispatch_key), 5); in expectCallsIncrementUnboxed()
70 void expectCallsDecrement(DispatchKey dispatch_key) { in expectCallsDecrement() argument
76 auto result = callOp(*op, dummyTensor(dispatch_key), 5); in expectCallsDecrement()
|
| D | test_helpers.h | 38 inline at::Tensor dummyTensor(c10::DispatchKey dispatch_key, bool requires_grad=false) {
39 return dummyTensor(c10::DispatchKeySet(dispatch_key), requires_grad);
64 inline void expectDoesntFindKernel(const char* op_name, c10::DispatchKey dispatch_key) { in expectDoesntFindKernel() argument
67 callOp(*op, dummyTensor(dispatch_key), 5); in expectDoesntFindKernel()
|
| /external/pytorch/c10/core/ |
| D | TensorOptions.h | 726 inline Layout dispatchKeyToLayout(DispatchKey dispatch_key) { in dispatchKeyToLayout() argument
727 switch (dispatch_key) { in dispatchKeyToLayout()
736 false, "Cannot map DispatchKey ", dispatch_key, " to a unique layout."); in dispatchKeyToLayout()
744 inline c10::DeviceType dispatchKeyToDeviceType(DispatchKey dispatch_key) { in dispatchKeyToDeviceType() argument
745 switch (dispatch_key) { in dispatchKeyToDeviceType()
766 dispatch_key, in dispatchKeyToDeviceType()
771 inline TensorOptions dispatchKeyToTensorOptions(DispatchKey dispatch_key) { in dispatchKeyToTensorOptions() argument
773 .layout(dispatchKeyToLayout(dispatch_key)) in dispatchKeyToTensorOptions()
774 .device(dispatchKeyToDeviceType(dispatch_key)); in dispatchKeyToTensorOptions()
|
| /external/pytorch/torch/csrc/autograd/ |
| D | python_legacy_variable.cpp | 68 auto dispatch_key = torch::tensors::get_default_dispatch_key(); in THPVariable_pynew() local
71 .device(dispatchKeyToDeviceType(dispatch_key)) in THPVariable_pynew()
72 .layout(dispatchKeyToLayout(dispatch_key)); in THPVariable_pynew()
|
| /external/pytorch/aten/src/ATen/test/ |
| D | lazy_tensor_test.cpp | 5 void LazyTensorTest(c10::DispatchKey dispatch_key, at::DeviceType device_type) { in LazyTensorTest() argument
8 dispatch_key, in LazyTensorTest()
|
| /external/pytorch/torch/_inductor/ |
| D | test_operators.py | 14 for dispatch_key in ("CPU", "CUDA", "Meta"):
15 _test_lib_impl.impl("realize", lambda x: x.clone(), dispatch_key)
|
| /external/executorch/examples/portable/custom_ops/ |
| D | custom_ops_1.py | 20 @impl(my_op_lib, "mul3", dispatch_key="CompositeExplicitAutograd")
31 @impl(my_op_lib, "mul3.out", dispatch_key="CompositeExplicitAutograd")
|
| /external/pytorch/tools/test/ |
| D | test_codegen.py | 343 dispatch_key=k,
459 dispatch_key = DispatchKey.CompositeExplicitAutograd
460 self.assertTrue(dispatch_key in self.backend_indices)
463 dispatch_key=dispatch_key,
467 index=self.backend_indices[dispatch_key],
|
| /external/pytorch/test/functorch/ |
| D | test_vmap_registrations.py | 236 dispatch_key: str, xfails: set, filter_func: typing.Callable = lambda reg: True
238 registrations = sorted(get_registrations_for_dispatch_key(dispatch_key))
|
| /external/pytorch/torchgen/api/ |
| D | ufunc.py | 28 def schema_kernel_name(func: FunctionSchema, dispatch_key: DispatchKey) -> str:
33 def kernel_name(g: NativeFunctionsGroup, dispatch_key: DispatchKey) -> str:
34 return schema_kernel_name(g.out.func, dispatch_key)
|