Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / _op_dispatch.py: 77%
617 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""_op_dispatch"""
16import atexit
17import glob
18import importlib
19import os
20import sys
21import warnings
22from contextvars import ContextVar
23from itertools import chain
24from typing import Any, Dict, FrozenSet, List, Optional
26import yaml
28from hyper_parallel.core.shard.ops.parallel_ops_register import get_distributed_op
29from hyper_parallel.core.dtensor.dtensor import DTensor
30from hyper_parallel.core.dtensor.random import OffsetBasedRNGTracker, is_rng_supported_mesh
31from hyper_parallel.platform import get_platform
32from hyper_parallel.platform.platform import PlatformType
34from hyper_parallel.core.tensor_parallel._ce_op_registry import is_loss_parallel_op, is_decomposed_ce_op
35from hyper_parallel.core.tensor_parallel.loss_parallel import is_loss_parallel_active
36from hyper_parallel.core.tensor_parallel.loss_parallel_ops_common import _is_shard_on_last_dim
38platform = get_platform()
39Tensor = platform.Tensor
42def _apply_shard_offset_to_rng_args(args, offset_incr):
43 """Apply per-shard offset increment to seed/offset tensors in MindSpore random op args.
45 MindSpore random ops (e.g. ``randn_like_``) receive ``(seed, offset)`` as
46 explicit int64 scalar tensors from ``default_generator._step()`` in the
47 Python wrapper *before* the C++ dispatch triggers ``__fallback__``. By the
48 time ``_dispatch_random_op`` is called, the kernel will use whatever
49 ``(seed, offset)`` values are in the args—it does **not** read the
50 generator again. This function finds the offset tensor and adds the
51 per-rank offset increment so each shard gets a unique random stream.
53 The (seed, offset) pair is identified as the last two consecutive int64
54 0-dim tensors in *args* (scanning from the end to skip trailing dtype /
55 device arguments).
57 Args:
58 args: The list of local args for the random op.
59 offset_incr (int): Per-shard offset increment.
61 Returns:
62 list: Modified args with the offset tensor adjusted.
63 """
64 int64_dtype = platform.tensor_dtype.int64
65 last_int64_idx = -1
66 for i in range(len(args) - 1, -1, -1):
67 arg = args[i]
68 if isinstance(arg, Tensor) and arg.dtype == int64_dtype and arg.ndim == 0:
69 if last_int64_idx == i + 1:
70 offset_idx = i + 1
71 new_args = list(args)
72 new_offset = int(new_args[offset_idx].item()) + offset_incr
73 new_args[offset_idx] = platform.tensor([new_offset], dtype=int64_dtype).reshape(())
74 return new_args
75 last_int64_idx = i
76 return args
78_dtensor_dispatch_disabled: ContextVar[bool] = ContextVar('_dtensor_dispatch_disabled', default=False)
79_no_skip_ops: ContextVar[FrozenSet[str]] = ContextVar('_no_skip_ops', default=frozenset())
82def get_no_skip_ops() -> FrozenSet[str]:
83 """Return the set of op names that are exempt from SkipDTensorDispatch."""
84 return _no_skip_ops.get()
87def get_dtensor_dispatch() -> bool:
88 """
89 Get the current DTensor dispatch status.
91 Returns:
92 bool: True if DTensor dispatch is enabled, False otherwise.
93 """
94 return not _dtensor_dispatch_disabled.get()
97class LayoutCacheKey:
98 """Immutable layout cache key."""
99 __slots__ = ('_tuple', '_hash')
101 def __init__(self, layout_ids: List[str]):
102 self._tuple = tuple(layout_ids)
103 self._hash = hash(self._tuple)
105 @classmethod
106 def from_cache_values(cls, cache_values: list) -> "LayoutCacheKey":
107 """Build a LayoutCacheKey from a cache_values list.
109 Args:
110 cache_values (list): Mixed list of Layout objects (with compact_str) and raw scalars.
112 Returns:
113 LayoutCacheKey: Immutable key derived from the string representation of each value.
114 """
115 key_values = []
116 for v in cache_values:
117 if hasattr(v, 'compact_str'):
118 key_values.append(str(v.compact_str))
119 else:
120 key_values.append(str(v))
121 return cls(key_values)
123 def __eq__(self, other):
124 if not isinstance(other, LayoutCacheKey):
125 return False
126 return self._tuple == other._tuple
128 def __hash__(self):
129 return self._hash
131 def __repr__(self):
132 return f"LayoutCacheKey({self._tuple})"
135class LayoutCacheManager:
136 """
137 Cache layout in infer layout.
139 A singleton class that manages layout caches for distributed operations.
140 It caches the inferred layouts and operation implementations to avoid
141 redundant computation during repeated calls with the same input layouts.
142 """
143 _instance = None
145 def __init__(self):
146 self.layout_cache: Dict[str, Dict[LayoutCacheKey, Any]] = {}
147 atexit.register(self.clear_cache)
149 @classmethod
150 def get_instance(cls) -> "LayoutCacheManager":
151 """
152 Get the singleton instance of LayoutCacheManager.
154 Returns:
155 LayoutCacheManager: The singleton instance.
156 """
157 if cls._instance is None:
158 cls._instance = LayoutCacheManager()
159 return cls._instance
161 def get_layout_cache(self) -> Dict[str, Dict[LayoutCacheKey, Any]]:
162 """
163 Get the layout cache dictionary.
165 Returns:
166 Dict[str, Dict[LayoutCacheKey, Any]]: The nested dictionary mapping
167 operation names to their layout caches.
168 """
169 return self.layout_cache
171 @staticmethod
172 def distributed_op(op_name: str) -> Any:
173 """
174 Get the distributed operation implementation by name.
176 Args:
177 op_name (str): The name of the distributed operation.
179 Returns:
180 Any: The distributed operation class or implementation.
181 """
182 op = get_distributed_op(op_name)
183 return op
185 def clear_cache(self) -> None:
186 """
187 Clear all cached layouts.
189 This method is automatically registered with atexit to ensure
190 cache is cleared when the program exits.
191 """
192 self.layout_cache.clear()
195class OpDispatcher:
196 """
197 OpDispatcher
198 """
200 # Whitelisted ops that mutate args[0]'s storage in place. The dispatch bypass
201 # must return the original DTensor self for these, not the unwrapped local
202 # result, or it demotes a DTensor accumulator to a plain Tensor and breaks the
203 # next op that adds a DTensor to it (e.g. grad-accumulation `loss += micro_loss`).
204 # Class-level so it stays available on instances built via __new__ (e.g. tests).
205 _INPLACE_BYPASS_OPS = frozenset(
206 {"InplaceAddExt", "InplaceSubExt", "InplaceMul", "InplaceDiv"})
208 # MindSpore random kernels that always mutate an existing tensor in place.
209 # Out-of-place random kernels belong in _random_ms_ops only, not here.
210 _RANDOM_INPLACE_MS_OPS = frozenset({
211 "InplaceBernoulliScalar",
212 "InplaceBernoulliTensor",
213 "InplaceNormal",
214 "InplaceRandom",
215 "InplaceUniform",
216 })
218 def __init__(self):
219 self._env_yaml_dir: Optional[str] = os.environ.get("HYPER_PARALLEL_OPS_YAML_DIR")
220 self._env_python_path: Optional[str] = os.environ.get("HYPER_PARALLEL_OPS_PYTHON_PATH")
221 # The following attributes are initialized in _setup_yaml_dir()
222 self.work_dir = "" # Initialized in _setup_yaml_dir()
223 self.yaml_dir = "" # Initialized in _setup_yaml_dir()
225 self._setup_paths_from_env()
227 self.layout_infer_ops = self.safe_load_yaml_from_dir()
228 self.whitelist = ["typeof", "DistCommIsend",
229 "DistCommIrecv", "DistCommBroadcast", "DistCommAllReduce", "DistCommAllGather", "DistCommBatchIsendIrecv",
230 "DistCommReduceScatter", "requires_grad_", "item", "__get__", "__set__", "register_hook",
231 "is_complex", "chunk", "__bool__", "__len__", "__format__", "dim",
232 "_has_compatible_shallow_copy_type", "is_floating_point", "is_contiguous"]
234 # Ops requiring args unpacking for layout inference (packed as prim, name, real_args).
235 self.unpack_ops = ["ScatterUpdate", "Mod", "GatherNd", "StopGradient"]
237 self._random_ops = {
238 "normal_", "uniform_", "bernoulli", "bernoulli_",
239 "native_dropout", "rand", "rand_like", "randn",
240 "randn_like", "randint_like", "kaiming_uniform_",
241 "multinomial",
242 }
243 # Only mint random op support
244 # MindSpore use the actual kernel name.
245 self._random_ms_ops = {
246 "BernoulliExt", "MultinomialExt",
247 "InplaceBernoulliScalar", "InplaceBernoulliTensor",
248 "InplaceNormal", "InplaceRandom", "InplaceUniform",
249 "NormalFloatFloat", "NormalFloatTensor", "NormalTensorFloat", "NormalTensorTensor",
250 "RandpermExt", "Randn", "RandLikeExt", "RandnLike", "RandInt", "RandIntLike", "RandExt",
251 "FuncDropoutExt", "UniformExt",
252 }
253 self._rng_tracker: Optional[OffsetBasedRNGTracker] = None
255 self._suffix_dispatch: Dict[str, str] = {
256 "WithShape": "_with_layout_infer_with_shape",
257 "Reshape": "_with_layout_infer_reshape",
258 "WithTupleExpand": "_with_layout_infer_with_tuple_expand",
259 "Slice": "_with_layout_infer_slice",
260 }
262 self._register_distributed_ops()
264 def _setup_paths_from_env(self):
265 """
266 Setup YAML directory and Python path from environment variables.
268 This method initializes the YAML directory and extends sys.path based on
269 environment variables HYPER_PARALLEL_OPS_YAML_DIR and HYPER_PARALLEL_OPS_PYTHON_PATH.
270 """
271 self._setup_yaml_dir(self._env_yaml_dir)
272 self._extend_sys_path(self._env_python_path)
274 def _setup_yaml_dir(self, env_yaml_dir: Optional[str]):
275 """
276 Feature: Configure yaml_dir/work_dir for OpDispatcher
277 Description: Resolve the YAML directory used to load distributed op definitions.
278 If env_yaml_dir is an absolute path, use it directly; otherwise treat it
279 as a path relative to the project work_dir. If env_yaml_dir is not set,
280 fall back to the default 'shard/ops/yaml' under work_dir.
281 Expectation: self.yaml_dir and self.work_dir are set to valid values used later by
282 safe_load_yaml_from_dir(); no functional behavior is changed.
283 """
284 if env_yaml_dir:
285 if os.path.isabs(env_yaml_dir):
286 self.yaml_dir = env_yaml_dir
287 self.work_dir = ""
288 else:
289 self.work_dir = os.path.normpath(
290 os.path.join(os.path.dirname(os.path.realpath(__file__)), "../")
291 )
292 self.yaml_dir = env_yaml_dir
293 else:
294 self.yaml_dir = "shard/ops/yaml"
295 self.work_dir = os.path.normpath(
296 os.path.join(os.path.dirname(os.path.realpath(__file__)), "../")
297 )
299 @staticmethod
300 def _extend_sys_path(env_python_path: Optional[str]):
301 if not env_python_path:
302 return
303 python_paths = env_python_path.split(":")
304 for path in python_paths:
305 if path and os.path.isdir(path) and path not in sys.path:
306 sys.path.append(path)
308 def _register_distributed_ops(self):
309 for op_name, config in self.layout_infer_ops.items():
310 self._register_single_distributed_op(op_name, config)
312 def _register_single_distributed_op(self, op_name: str, config: dict):
313 """
314 Feature: Register a single distributed op implementation
315 Description: Import the distributed op class specified by config and instantiate it
316 with op_name to trigger registration in the distributed op registry.
317 Prefer 'distributed_op_module' when provided; otherwise import from
318 built-in module prefix 'hyper_parallel.core.shard.ops.' plus
319 'distributed_op_file'. If import fails and an external python path is
320 provided via env, fall back to importing 'distributed_op_file' directly.
321 Expectation: The distributed op class is imported and instantiated successfully,
322 or the original import error is raised; no functional behavior is changed.
323 """
324 class_name = config["distributed_op_class"]
326 if "distributed_op_module" in config:
327 module_name = config["distributed_op_module"]
328 module = importlib.import_module(module_name)
329 op_class = getattr(module, class_name)
330 _ = op_class(op_name)
331 return
333 module_file = config["distributed_op_file"]
334 try:
335 module_name = "hyper_parallel.core.shard.ops." + module_file
336 module = importlib.import_module(module_name)
337 op_class = getattr(module, class_name)
338 _ = op_class(op_name)
339 except (ModuleNotFoundError, ImportError):
340 if self._env_python_path:
341 module = importlib.import_module(module_file)
342 op_class = getattr(module, class_name)
343 _ = op_class(op_name)
344 else:
345 raise
347 @staticmethod
348 def _process_args_and_kwargs(
349 args, kwargs
350 ) -> tuple[list, list, list, dict, list]:
351 """_process_args_and_kwargs"""
352 input_layouts = []
353 extra_args = []
354 input_args = []
355 input_kwargs = kwargs.copy()
356 cache_key_values = []
358 for arg in args:
359 if arg is None:
360 input_layouts.append(None)
361 input_args.append(arg)
362 continue
364 if not hasattr(arg, "_layout"):
365 id_str = "scalar"
366 if not isinstance(arg, Tensor):
367 id_str = str(arg)
368 cache_key_values.append(id_str)
369 extra_args.append(arg)
370 input_layouts.append(None)
371 input_args.append(arg)
372 else:
373 layout = arg.layout
374 layout_id = layout.compact_str
375 cache_key_values.append(str(layout_id))
376 input_layouts.append(layout)
377 if isinstance(arg, DTensor):
378 input_args.append(arg.to_local())
379 else:
380 input_args.append(arg)
382 for k, val in kwargs.items():
383 if val is None:
384 input_layouts.append(None)
385 continue
386 if not hasattr(val, "_layout"):
387 id_str = "scalar"
388 if not isinstance(val, Tensor):
389 id_str = str(val)
390 cache_key_values.append(id_str)
391 extra_args.append(val)
392 input_layouts.append(None)
393 else:
394 layout = val.layout
395 layout_id = layout.compact_str
396 cache_key_values.append(str(layout_id))
397 input_layouts.append(layout)
398 if isinstance(val, DTensor):
399 input_kwargs[k] = val.to_local()
401 return input_layouts, extra_args, input_args, input_kwargs, cache_key_values
403 @staticmethod
404 def _with_layout_infer(func: callable, *args, _packed_call=None, **kwargs) -> Tensor:
405 """_with_layout_infer
407 NOTE: aclop packed args are already normalized upstream by
408 ``_dispatch_layout_infer`` via ``_normalize_aclop_args``. The
409 ``_packed_call`` kwarg carries the ``(prim, name)`` tuple when the
410 op uses aclop format, or ``None`` otherwise.
411 """
412 func_name = platform.get_op_name(func)
414 input_layouts, extra_args, input_args, input_kwargs, cache_values = \
415 OpDispatcher._process_args_and_kwargs(args, kwargs)
416 cache_key = LayoutCacheKey(cache_values)
418 output_layout, op_impl, _ = OpDispatcher._lookup_or_infer_layout(
419 func, func_name, cache_key, input_layouts, extra_args
420 )
422 op_impl = func if op_impl is None else op_impl
423 py_output = OpDispatcher._call_op_impl(op_impl, _packed_call, input_args, input_kwargs)
424 return OpDispatcher._pack_infer_output(py_output, output_layout)
426 @staticmethod
427 def _extract_single_arg_layout(expanded_args, kwargs_value):
428 """Helper to extract layout and cache info for a single argument."""
429 cache_key_values = []
430 input_layouts = []
431 extra_args = []
433 for arg in chain(expanded_args, kwargs_value):
434 if arg is None:
435 input_layouts.append(None)
436 continue
438 if not hasattr(arg, "_layout"):
439 id_str = "scalar" if isinstance(arg, Tensor) else str(arg)
440 cache_key_values.append(id_str)
441 extra_args.append(arg)
442 input_layouts.append(None)
443 else:
444 layout = arg.layout
445 cache_key_values.append(str(layout.compact_str))
446 input_layouts.append(layout)
447 return cache_key_values, input_layouts, extra_args
449 @staticmethod
450 def _pack_infer_output(py_output, output_layout):
451 """Helper to pack py_output into DTensors using output_layout."""
452 if isinstance(py_output, (tuple, list)):
453 if not isinstance(output_layout, (tuple, list)):
454 raise RuntimeError("Output is a tuple but layout is not")
455 if len(py_output) != len(output_layout):
456 raise RuntimeError(f"Output tuple size ({len(py_output)}) "
457 f"does not match layout tuple size ({len(output_layout)})")
459 return tuple(
460 DTensor.from_local(item, layout.mesh, layout.alias_placements)
461 for item, layout in zip(py_output, output_layout)
462 )
464 return DTensor.from_local(py_output, output_layout.mesh, output_layout.alias_placements)
466 @staticmethod
467 def _expand_tuple_args(args: tuple) -> tuple:
468 """Expand tuple/list arguments for tuple element-wise ops.
470 Returns:
471 (expanded_args, input_args): Flat list of DTensors for layout extraction
472 and list of local-tensor tuples/values for kernel invocation.
473 """
474 expanded_args = []
475 input_args = []
476 for arg in args:
477 if isinstance(arg, (tuple, list)):
478 expanded_args.extend(arg)
479 # pylint: disable=R1728
480 input_args.append(tuple(
481 item.to_local() if hasattr(item, "_layout") else item for item in arg
482 ))
483 else:
484 expanded_args.append(arg)
485 input_args.append(arg.to_local() if isinstance(arg, DTensor) else arg)
486 return expanded_args, input_args
488 @staticmethod
489 def _lookup_or_infer_layout(func, func_name, cache_key, input_layouts, extra_args):
490 """Look up cached layout or compute via distributed op.
492 Returns:
493 (output_layout, op_impl, distribute_op)
494 """
495 cache_manager = LayoutCacheManager.get_instance()
496 layout_cache = cache_manager.get_layout_cache()
497 if func_name not in layout_cache:
498 layout_cache[func_name] = {}
500 op_layout_cache = layout_cache[func_name]
501 distribute_op = cache_manager.distributed_op(func_name)
503 if cache_key in op_layout_cache:
504 output_layout, op_impl = op_layout_cache[cache_key]
505 else:
506 output_layout = distribute_op.infer_layout(input_layouts, extra_args)
507 op_impl = distribute_op.get_expand_impl(func, output_layout, input_layouts, extra_args)
508 op_layout_cache[cache_key] = (output_layout, op_impl)
510 return output_layout, op_impl, distribute_op
512 @staticmethod
513 def _prepare_tuple_expand_args(args, kwargs):
514 """Prepare arguments for tuple-expand layout inference.
516 Returns:
517 (input_args, input_kwargs, input_layouts, extra_args, cache_key)
518 """
519 expanded_args, input_args = OpDispatcher._expand_tuple_args(args)
520 input_kwargs = {
521 key: value.to_local() if isinstance(value, DTensor) else value
522 for key, value in kwargs.items()
523 }
524 cache_values, input_layouts, extra_args = OpDispatcher._extract_single_arg_layout(
525 expanded_args, kwargs.values()
526 )
527 return input_args, input_kwargs, input_layouts, extra_args, LayoutCacheKey(cache_values)
529 def _with_layout_infer_with_tuple_expand(self, func: callable, *args, _packed_call=None, **kwargs) -> Tensor:
530 """_with_layout_infer_with_tuple_expand
532 NOTE: aclop packed args are already normalized upstream by
533 ``_dispatch_layout_infer`` via ``_normalize_aclop_args``. The
534 ``_packed_call`` kwarg carries the ``(prim, name)`` tuple when the
535 op uses aclop format, or ``None`` otherwise.
536 """
537 input_args, input_kwargs, input_layouts, extra_args, cache_key = (
538 self._prepare_tuple_expand_args(args, kwargs)
539 )
540 func_name = platform.get_op_name(func)
542 output_layout, op_impl, distribute_op = OpDispatcher._lookup_or_infer_layout(
543 func, func_name, cache_key, input_layouts, extra_args
544 )
546 op_impl = func if op_impl is None else op_impl
547 py_output = OpDispatcher._call_op_impl(op_impl, _packed_call, input_args, input_kwargs)
548 return distribute_op.wrap_output(py_output, output_layout)
550 @staticmethod
551 def _with_layout_infer_reshape(func: callable, *args) -> Tensor:
552 """_with_layout_infer_reshape"""
553 input_tensor = args[0]
555 # PyTorch reshape signature: reshape(*shape)
556 # Can be called as:
557 # reshape(-1, 1024) -> args = (tensor, -1, 1024)
558 # reshape((-1, 1024)) -> args = (tensor, (-1, 1024))
559 # reshape([4, 16, 1024]) -> args = (tensor, [4, 16, 1024])
560 if len(args) > 2:
561 # Multiple int arguments: reshape(-1, 1024)
562 shape = args[1:]
563 else:
564 # Single argument: reshape((4, 16, 1024)) or reshape(4)
565 shape = args[1]
567 layout = input_tensor.layout
568 input_layouts = [layout]
570 extra_args = [shape, input_tensor.shape]
572 cache_key_values = [str(layout.compact_str), str(shape), str(input_tensor.shape)]
573 cache_key = LayoutCacheKey(cache_key_values)
575 cache_manager = LayoutCacheManager.get_instance()
576 layout_cache = cache_manager.get_layout_cache()
577 func_name = platform.get_op_name(func)
578 if func_name not in layout_cache:
579 layout_cache[func_name] = {}
581 op_layout_cache = layout_cache[func_name]
583 distribute_op = cache_manager.distributed_op(func_name)
584 if cache_key in op_layout_cache:
585 infer_output, op_impl = op_layout_cache[cache_key]
586 else:
587 all_args = (input_layouts, extra_args)
588 infer_output = distribute_op.infer_layout(*all_args)
589 op_impl = distribute_op.get_expand_impl(func, infer_output, input_layouts, extra_args)
590 op_layout_cache[cache_key] = (infer_output, op_impl)
592 infer_output_tuple = infer_output
593 local_shape = infer_output_tuple[1]
595 if op_impl is None:
596 op_impl = func
598 py_output = op_impl(input_tensor.to_local(), local_shape)
600 return DTensor.from_local(py_output, infer_output_tuple[0].mesh, infer_output_tuple[0].alias_placements)
602 @staticmethod
603 def _process_args_and_kwargs_with_shape(args, kwargs):
604 """Process args and kwargs with input shapes for WithShape suffix operators.
606 Args:
607 args: Positional arguments from dispatch.
608 kwargs: Keyword arguments from dispatch.
610 Returns:
611 tuple: (input_layouts, input_shapes, extra_args, input_args, input_kwargs, cache_key_values)
612 """
613 input_layouts = []
614 extra_args = []
615 input_shapes = []
616 input_args = []
617 input_kwargs = kwargs.copy()
618 cache_key_values = []
620 for arg in args:
621 if arg is None:
622 input_layouts.append(None)
623 input_shapes.append(None)
624 input_args.append(arg)
625 continue
627 if not hasattr(arg, "_layout"):
628 id_str = "scalar"
629 if not isinstance(arg, Tensor):
630 id_str = str(arg)
631 cache_key_values.append(id_str)
632 extra_args.append(arg)
633 input_layouts.append(None)
634 input_args.append(arg)
635 else:
636 layout = arg.layout
637 layout_id = layout.compact_str
638 cache_key_values.append(str(layout_id))
639 input_layouts.append(layout)
640 if isinstance(arg, DTensor):
641 input_args.append(arg.to_local())
642 else:
643 input_args.append(arg)
645 if not hasattr(arg, "shape"):
646 input_shapes.append(None)
647 else:
648 input_shape = arg.shape
649 input_shapes.append(input_shape)
650 cache_key_values.append(str(input_shape))
652 for k, val in kwargs.items():
653 if val is None:
654 input_layouts.append(None)
655 continue
656 if not hasattr(val, "_layout"):
657 id_str = "scalar"
658 if not isinstance(val, Tensor):
659 id_str = str(val)
660 cache_key_values.append(id_str)
661 extra_args.append(val)
662 input_layouts.append(None)
663 else:
664 layout = val.layout
665 layout_id = layout.compact_str
666 cache_key_values.append(str(layout_id))
667 input_layouts.append(layout)
668 if isinstance(val, DTensor):
669 input_kwargs[k] = val.to_local()
671 if not hasattr(val, "shape"):
672 input_shapes.append(None)
673 else:
674 input_shape = val.shape
675 input_shapes.append(input_shape)
676 cache_key_values.append(str(input_shape))
678 return input_layouts, input_shapes, extra_args, input_args, input_kwargs, cache_key_values
680 @staticmethod
681 def _prepare_with_shape_args(args, kwargs):
682 """Prepare arguments and cache key for WithShape operators.
684 Returns:
685 (layouts, extra_args, local_args, local_kwargs, cache_key)
686 """
687 processed = OpDispatcher._process_args_and_kwargs_with_shape(args, kwargs)
688 layouts, shapes, extra_args, local_args, local_kwargs, cache_values = processed
689 extra_args.append(shapes)
690 return layouts, extra_args, local_args, local_kwargs, LayoutCacheKey(cache_values)
692 @staticmethod
693 def _with_layout_infer_with_shape(func: callable, *args, _packed_call=None, **kwargs) -> Tensor:
694 """_with_layout_infer_with_shape
696 NOTE: aclop packed args are already normalized upstream by
697 ``_dispatch_layout_infer`` via ``_normalize_aclop_args``. The
698 ``_packed_call`` kwarg carries the ``(prim, name)`` tuple when the
699 op uses aclop format, or ``None`` otherwise.
700 """
701 prepared = OpDispatcher._prepare_with_shape_args(args, kwargs)
702 layouts, extra_args, local_args, local_kwargs, cache_key = prepared
703 func_name = platform.get_op_name(func)
705 infer_result = OpDispatcher._lookup_or_infer_layout(
706 func, func_name, cache_key, layouts, extra_args
707 )
708 output_layout, op_impl = infer_result[:2]
710 op_impl = func if op_impl is None else op_impl
711 py_output = OpDispatcher._call_op_impl(op_impl, _packed_call, local_args, local_kwargs)
712 return OpDispatcher._pack_infer_output(py_output, output_layout)
714 def _with_layout_infer_slice(self, func: callable, *args) -> Tensor:
715 """_with_layout_infer_slice"""
716 input_tensor = args[0]
717 begin = args[1]
718 end = args[2]
720 # input layout
721 input_layouts = []
723 layout = input_tensor.layout
724 global_shape = input_tensor.shape
725 input_layouts.append(layout)
726 layout_id = layout.compact_str
728 extra_args = []
729 extra_args.append(begin)
730 extra_args.append(end)
731 extra_args.append(global_shape)
732 cache_key_values = [str(layout_id), str(begin), str(end), str(global_shape)]
733 cache_key = LayoutCacheKey(cache_key_values)
735 cache_manager = LayoutCacheManager.get_instance()
736 layout_cache = cache_manager.get_layout_cache()
737 func_name = platform.get_op_name(func)
738 if func_name not in layout_cache:
739 layout_cache[func_name] = {}
741 op_layout_cache = layout_cache[func_name]
743 distribute_op = cache_manager.distributed_op(func_name)
744 if cache_key in op_layout_cache:
745 infer_output, op_impl = op_layout_cache[cache_key]
746 else:
747 all_args = (input_layouts, extra_args)
748 infer_output = distribute_op.infer_layout(*all_args)
749 op_impl = distribute_op.get_expand_impl(func, infer_output, input_layouts, extra_args)
750 op_layout_cache[cache_key] = (infer_output, op_impl)
752 infer_output_tuple = infer_output
753 new_begin = infer_output_tuple[1]
754 new_end = infer_output_tuple[2]
756 if op_impl is None:
757 op_impl = func
759 py_output = op_impl(input_tensor.to_local(), new_begin, new_end)
761 return DTensor.from_local(py_output, infer_output_tuple[0].mesh, infer_output_tuple[0].alias_placements)
763 @staticmethod
764 def _merge_default(config: dict):
765 """Apply __default__ values to all ops in this YAML file."""
766 if "__default__" not in config:
767 return config
769 default_cfg = config["__default__"]
770 merged = {}
772 for op_name, op_cfg in config.items():
773 if op_name == "__default__":
774 continue
776 new_cfg = default_cfg.copy()
777 new_cfg.update(op_cfg)
778 merged[op_name] = new_cfg
780 return merged
782 def safe_load_yaml_from_dir(self) -> dict:
783 """
784 Load yaml dictionary from directory.
786 Returns:
787 dict: Merged dictionary of all operator configurations loaded from YAML files.
788 """
789 yaml_dict = {}
790 yaml_path = os.path.join(self.work_dir, self.yaml_dir) if self.work_dir else self.yaml_dir
791 if not os.path.isdir(yaml_path):
792 raise ValueError(f"Invalid yaml directory path: {yaml_path}")
794 for yaml_file_path in glob.glob(os.path.join(yaml_path, '*.yaml')):
795 with open(yaml_file_path, 'r', encoding="utf-8") as f:
796 yaml_data = yaml.safe_load(f)
798 yaml_data = OpDispatcher._merge_default(yaml_data)
799 for name, data in yaml_data.items():
800 if name in yaml_dict:
801 raise ValueError(f"Duplicate yaml object with name '{name}'.")
802 yaml_dict[name] = data
804 return yaml_dict
806 def _dispatch_random_op(self, op_name: str, op_call: callable, args, kwargs):
807 """Handle dispatch for random ops that operate on DTensors."""
808 first_arg = next(
809 (x for x in chain(args, kwargs.values()) if isinstance(x, DTensor)),
810 None,
811 )
812 # Fall back to the default op if no DTensor is found.
813 if first_arg is None:
814 return op_call(*args, **kwargs)
816 local_args = [arg.to_local() if isinstance(arg, DTensor) else arg for arg in args]
817 local_kwargs = {k: v.to_local() if isinstance(v, DTensor) else v for k, v in kwargs.items()}
818 first_local_arg = first_arg.to_local()
820 if self._rng_tracker is None and is_rng_supported_mesh(first_arg.device_mesh):
821 self._rng_tracker = OffsetBasedRNGTracker()
823 maybe_user_generator = local_kwargs.pop("generator", None)
824 if (
825 self._rng_tracker is not None
826 and not first_local_arg.is_meta
827 and self._rng_tracker.distribute_region_enabled
828 ):
829 # pylint: disable=W0212
830 with self._rng_tracker._distribute_region(
831 device_mesh=first_arg.device_mesh,
832 placements=first_arg.placements,
833 global_shape=first_arg.shape,
834 generator=maybe_user_generator,
835 ):
836 # MindSpore random ops (e.g. mint.randn_like) extract (seed, offset)
837 # from default_generator._step() in the Python wrapper *before* the
838 # C++ dispatch triggers __fallback__. The callback reuses these
839 # pre-fetched tensor args, so set_rng_state inside _distribute_region
840 # has no effect on the kernel. Fix: apply the per-shard offset
841 # increment directly to the offset tensor in the args.
842 if platform.platform_type == PlatformType.MINDSPORE:
843 offset_incr = self._rng_tracker.compute_offset_incr(
844 first_arg.device_mesh, first_arg.placements, first_arg.shape,
845 )
846 local_args = _apply_shard_offset_to_rng_args(local_args, offset_incr)
847 local_results = op_call(*local_args, **local_kwargs)
848 else:
849 if maybe_user_generator is not None:
850 local_kwargs["generator"] = maybe_user_generator
851 local_results = op_call(*local_args, **local_kwargs)
853 return self._wrap_random_result(op_name, local_results, first_arg, args, kwargs)
855 @staticmethod
856 def _func_dropout_ext_inplace(args, kwargs) -> bool:
857 """Return True when FuncDropoutExt is invoked with inplace=True."""
858 # Kernel signature: (input, p, training, inplace, seed, offset).
859 if len(args) >= 4:
860 return bool(args[3])
861 return bool(kwargs.get("inplace", False))
863 @staticmethod
864 def _random_op_returns_self(op_name: str, args, kwargs) -> bool:
865 """Return True when a random op mutates an existing DTensor in place."""
866 if op_name in OpDispatcher._RANDOM_INPLACE_MS_OPS:
867 return True
868 if op_name == "FuncDropoutExt":
869 return OpDispatcher._func_dropout_ext_inplace(args, kwargs)
870 # Torch random inplace ops follow the ATen '_' suffix convention.
871 return op_name.endswith('_')
873 @staticmethod
874 def _wrap_random_result(op_name, local_results, first_arg, args, kwargs):
875 """Wrap a random op's local result(s) back into DTensor(s).
877 In-place ops return the input DTensor itself. Torch random inplace ops use
878 the ATen '_' suffix; MindSpore inplace random kernels are listed in
879 ``_RANDOM_INPLACE_MS_OPS``. ``FuncDropoutExt`` is handled separately
880 because the same kernel serves both modes via its ``inplace`` argument.
881 """
882 if OpDispatcher._random_op_returns_self(op_name, args, kwargs):
883 return first_arg
884 mesh = first_arg.device_mesh
885 placements = first_arg.layout.alias_placements
886 # Some ops return tuple/list, e.g. native_dropout returns (output, mask).
887 if isinstance(local_results, (tuple, list)):
888 return tuple(
889 DTensor.from_local(r, mesh, placements) if isinstance(r, Tensor) else r
890 for r in local_results
891 )
892 if isinstance(local_results, Tensor):
893 return DTensor.from_local(local_results, mesh, placements)
894 # Fallback: return as-is for non-Tensor results (currently unreachable with existing _random_ops).
895 return local_results
897 @staticmethod
898 def _unwrap_value(value: object) -> object:
899 """Replace DTensor with its local tensor; pass scalars and plain tensors through.
901 Args:
902 value (object): A single argument value from an op call.
904 Returns:
905 object: The local tensor if value is a DTensor, otherwise value unchanged.
906 """
907 if isinstance(value, DTensor):
908 return value.to_local()
909 if isinstance(value, tuple):
910 return tuple(OpDispatcher._unwrap_value(e) for e in value)
911 if isinstance(value, list):
912 return [OpDispatcher._unwrap_value(e) for e in value]
913 return value
915 @staticmethod
916 def _unwrap_args(args: tuple) -> list:
917 """Strip DTensor wrappers from args, preserving tuple/list container structure.
919 Args:
920 args: Op call positional arguments, may contain DTensor instances.
922 Returns:
923 List of args with DTensor replaced by their local tensors.
924 """
925 return [OpDispatcher._unwrap_value(arg) for arg in args]
927 @staticmethod
928 def _unwrap_kwargs(kwargs: dict) -> dict:
929 """Strip DTensor wrappers from kwargs values, preserving tuple/list container structure.
931 Args:
932 kwargs: Op call keyword arguments, values may contain DTensor instances.
934 Returns:
935 Dict of kwargs with DTensor values replaced by their local tensors.
936 """
937 return {k: OpDispatcher._unwrap_value(v) for k, v in kwargs.items()}
939 @staticmethod
940 def _gather_dtensors_to_full(args: tuple, kwargs: dict) -> tuple:
941 """Gather all DTensor arguments to full tensors for fallback execution.
943 Used when an operator has no parallel layout implementation. All DTensor
944 arguments are gathered to full tensors before calling the standard operator.
946 Args:
947 args: Op call positional arguments, may contain DTensor instances.
948 kwargs: Op call keyword arguments, may contain DTensor instances.
950 Returns:
951 Tuple of (unwrapped_args, unwrapped_kwargs) with DTensor values
952 replaced by their full tensor representations.
954 Warning:
955 This fallback performs all-gather which may consume significant memory.
956 Operators without layout implementations should be registered properly.
957 """
958 def gather(value: object) -> object:
959 if isinstance(value, DTensor):
960 return value.full_tensor()
961 if isinstance(value, tuple):
962 return tuple(gather(e) for e in value)
963 if isinstance(value, list):
964 return [gather(e) for e in value]
965 return value
967 gathered_args = [gather(arg) for arg in args]
968 gathered_kwargs = {k: gather(v) for k, v in kwargs.items()}
970 warnings.warn(
971 "Operator has no distributed layout implementation. "
972 "Falling back to all-gather which may consume significant memory. "
973 "Consider registering a proper distributed operator.",
974 UserWarning,
975 stacklevel=4
976 )
978 return gathered_args, gathered_kwargs
980 def _should_bypass_dispatch(self, op_name: str) -> bool:
981 """Return True if the op should bypass DTensor dispatch and run locally.
983 Args:
984 op_name: Canonical operator name from platform.get_op_name().
986 Returns:
987 True when the op is whitelisted or DTensor dispatch is globally disabled.
988 """
989 skip_dispatch = get_dtensor_dispatch() is False and op_name not in get_no_skip_ops()
990 return op_name in self.whitelist or op_name in self._INPLACE_BYPASS_OPS or skip_dispatch
992 def _should_dispatch_loss_parallel(self, op_name: str) -> bool:
993 """Check if should dispatch through loss_parallel path.
995 Args:
996 op_name: Canonical operator name from platform.get_op_name().
998 Returns:
999 True when in loss_parallel context and op is a CE entry point.
1000 """
1001 return is_loss_parallel_active() and is_loss_parallel_op(op_name)
1003 def _check_decomposed_ce_op_in_loss_parallel(self, op_name: str, args: tuple, kwargs: dict):
1004 """Check if decomposed CE ops are called in loss_parallel context.
1006 Args:
1007 op_name: Canonical operator name.
1008 args: Positional arguments for op_call.
1009 kwargs: Keyword arguments for op_call.
1011 Raises:
1012 ValueError: If decomposed CE op is called in loss_parallel context
1013 with vocab-sharded DTensor input.
1014 """
1015 if not is_loss_parallel_active() or not is_decomposed_ce_op(op_name):
1016 return
1018 has_vocab_sharded_dtensor = False
1019 for arg in args:
1020 if isinstance(arg, DTensor) and _is_shard_on_last_dim(arg):
1021 has_vocab_sharded_dtensor = True
1022 break
1023 if not has_vocab_sharded_dtensor:
1024 for val in kwargs.values():
1025 if isinstance(val, DTensor) and _is_shard_on_last_dim(val):
1026 has_vocab_sharded_dtensor = True
1027 break
1029 if has_vocab_sharded_dtensor:
1030 raise ValueError(
1031 f"Operator '{op_name}' is a decomposed component of cross_entropy and should not be called "
1032 f"directly within loss_parallel() context. Use F.cross_entropy(logits, targets) instead. "
1033 f"For example, replace:\n"
1034 f" with loss_parallel():\n"
1035 f" log_probs = F.log_softmax(logits, dim=-1)\n"
1036 f" loss = F.nll_loss(log_probs, targets)\n"
1037 f"with:\n"
1038 f" with loss_parallel():\n"
1039 f" loss = F.cross_entropy(logits, targets)"
1040 )
1042 def _dispatch_loss_parallel(self, op_call: callable, args: tuple, kwargs: dict):
1043 """Dispatch cross_entropy through the loss_parallel distributed kernel.
1045 Args:
1046 op_call: The raw operator callable.
1047 args: Positional arguments for op_call.
1048 kwargs: Keyword arguments for op_call.
1050 Returns:
1051 Result of the distributed cross_entropy computation.
1052 """
1053 if platform.platform_type == PlatformType.PYTORCH:
1054 # pylint: disable=C0415
1055 from hyper_parallel.platform.torch.loss_parallel_ops import distributed_cross_entropy_from_op_call
1056 elif platform.platform_type == PlatformType.MINDSPORE:
1057 # pylint: disable=C0415
1058 from hyper_parallel.platform.mindspore.loss_parallel_ops import distributed_cross_entropy_from_op_call
1059 else:
1060 raise RuntimeError(f"Unsupported platform for loss_parallel: {platform.platform_type}")
1061 return distributed_cross_entropy_from_op_call(op_call, args, kwargs)
1063 def _check_ce_op_without_loss_parallel_context(self, op_name: str, args: tuple):
1064 """Check if CE op is called with Shard(-1) DTensor outside loss_parallel context.
1066 Args:
1067 op_name: Canonical operator name.
1068 args: Positional arguments for op_call.
1070 Raises:
1071 ValueError: If CE op is called with Shard(-1) logits outside loss_parallel context.
1072 """
1073 if is_loss_parallel_active() or not is_loss_parallel_op(op_name):
1074 return
1076 if len(args) == 0 or not isinstance(args[0], DTensor):
1077 return
1079 logits = args[0]
1080 if _is_shard_on_last_dim(logits):
1081 raise ValueError(
1082 f"Operator '{op_name}' requires loss_parallel context when input logits are "
1083 f"sharded on the vocabulary dimension (Shard(-1)). Please wrap your forward "
1084 f"and backward pass with loss_parallel():\n"
1085 f" with loss_parallel():\n"
1086 f" loss = F.cross_entropy(logits, targets)\n"
1087 f" loss.backward()\n"
1088 f"If you intentionally want to gather all shards to compute cross_entropy "
1089 f"(not recommended for large vocabulary), use logits.full_tensor() explicitly."
1090 )
1092 @staticmethod
1093 def _normalize_aclop_args(op_name: str, unpack_ops: list, args: tuple) -> tuple:
1094 """
1095 Normalize aclop-packed arguments for MindSpore backend operators.
1097 NOTE: This handles MindSpore aclop operators whose kernel signature packs
1098 arguments as ``(prim, op_name_str, (real_arg0, real_arg1, ...))``. The
1099 ``prim`` and ``op_name_str`` are preserved as ``packed_call`` for the
1100 final kernel invocation, while the real tensor arguments are extracted
1101 for layout inference and preprocessing.
1103 **aclop is planned for deprecation.** Once aclop is fully removed, this
1104 normalization and the associated ``unpack_ops`` list can be deleted.
1106 Args:
1107 op_name (str): Canonical operator name.
1108 unpack_ops (list): List of op names that may use aclop packed format.
1109 args (tuple): Raw positional arguments from the op call.
1111 Returns:
1112 tuple: ``(packed_call, normalized_args)``
1113 - **packed_call**: ``(prim, op_name_str)`` tuple for kernel
1114 invocation, or ``None`` if no unpacking was performed.
1115 - **normalized_args**: The real tensor arguments (unpacked if
1116 the packed format was detected, otherwise the original args).
1117 """
1118 if OpDispatcher._is_aclop_packed(op_name, unpack_ops, args):
1119 return (args[0], args[1]), tuple(args[2])
1120 return None, args
1122 @staticmethod
1123 def _is_aclop_packed(op_name: str, unpack_ops: list, args: tuple) -> bool:
1124 """Check if arguments use aclop packed format."""
1125 return (
1126 op_name in unpack_ops
1127 and len(args) == 3
1128 and isinstance(args[1], str)
1129 and isinstance(args[2], (tuple, list))
1130 )
1132 @staticmethod
1133 def _call_op_impl(op_impl: callable, packed_call, args, kwargs: dict):
1134 """Invoke *op_impl* with optional aclop packed-call wrapping.
1136 When *packed_call* is not ``None`` the MindSpore aclop kernel expects
1137 ``(prim, op_name, (arg0, arg1, ...))``. Otherwise *args* are spread
1138 as positional arguments in the usual way.
1140 Args:
1141 op_impl: The op implementation callable.
1142 packed_call: ``(prim, op_name)`` tuple or ``None``.
1143 args: Local tensor arguments (list or tuple).
1144 kwargs: Keyword arguments dict.
1146 Returns:
1147 Result of the *op_impl* invocation.
1148 """
1149 if packed_call is not None:
1150 return op_impl(packed_call[0], packed_call[1], tuple(args), **kwargs)
1151 return op_impl(*args, **kwargs)
1153 def _dispatch_layout_infer(
1154 self, op_name: str, op_call: callable, args: tuple, kwargs: dict
1155 ):
1156 """Dispatch an op through the layout-inference path.
1158 Args:
1159 op_name: Canonical operator name.
1160 op_call: The raw operator callable.
1161 args: Positional arguments for op_call.
1162 kwargs: Keyword arguments for op_call.
1164 Returns:
1165 Result of the layout-infer dispatch.
1167 Raises:
1168 RuntimeError: If op_name is not registered or has an unknown suffix.
1169 """
1170 if op_name not in self.layout_infer_ops:
1171 has_dtensor = any(isinstance(arg, DTensor) for arg in args)
1172 has_dtensor = has_dtensor or any(isinstance(v, DTensor) for v in kwargs.values())
1173 if has_dtensor:
1174 self._check_ce_op_without_loss_parallel_context(op_name, args)
1176 if not is_loss_parallel_op(op_name):
1177 raise RuntimeError(
1178 f"Operator {op_name} does not contain parallel layout infer func. "
1179 f"DTensor dispatch requires explicit layout inference registration. "
1180 f"Please register a distributed operator for '{op_name}' or use local tensors."
1181 )
1183 gathered_args, gathered_kwargs = self._gather_dtensors_to_full(args, kwargs)
1185 # Special handling for cross_entropy with 3D logits (only when NOT in loss_parallel context)
1186 # PyTorch expects: logits [N, C], targets [N]
1187 # But LLM forward returns: logits [batch, seq, vocab], targets [batch, seq]
1188 # Note: nll_loss input is log_probs, typically already 2D, so we only reshape for cross_entropy
1189 if op_name == "cross_entropy" and len(gathered_args) >= 2:
1190 logits = gathered_args[0]
1191 targets = gathered_args[1]
1192 if isinstance(logits, Tensor) and isinstance(targets, Tensor):
1193 if logits.ndim > 2 and targets.ndim > 1 and targets.ndim == logits.ndim - 1:
1194 vocab_size = logits.shape[-1]
1195 gathered_args[0] = logits.reshape(-1, vocab_size)
1196 gathered_args[1] = targets.reshape(-1)
1198 return op_call(*gathered_args, **gathered_kwargs)
1199 raise RuntimeError(f"Operator {op_name} does not contain parallel layout infer func.")
1201 cache_manager = LayoutCacheManager.get_instance()
1202 distribute_op = cache_manager.distributed_op(op_name)
1204 # Normalize aclop-packed args before any per-op processing.
1205 # This allows all distributed op classes (ElementWiseDistributedOp,
1206 # GatherNdDistributedOp, etc.) to receive clean unpacked args,
1207 # avoiding duplicated unpack logic in each class's preprocess.
1208 packed_call, args = self._normalize_aclop_args(op_name, getattr(self, 'unpack_ops', []), args)
1210 result = distribute_op.preprocess(args, kwargs)
1211 if result is not None:
1212 return self._dispatch_new(op_call, distribute_op, packed_call, result)
1214 suffix = self.layout_infer_ops[op_name].get('infer_layout_suffix', '')
1215 if not suffix:
1216 return self._with_layout_infer(op_call, *args, _packed_call=packed_call, **kwargs)
1218 if suffix == 'WithShape':
1219 return self._with_layout_infer_with_shape(op_call, *args, _packed_call=packed_call, **kwargs)
1221 if suffix == 'WithTupleExpand':
1222 return self._with_layout_infer_with_tuple_expand(op_call, *args, _packed_call=packed_call, **kwargs)
1224 handler_name = self._suffix_dispatch.get(suffix)
1225 if handler_name is None:
1226 raise RuntimeError(f"Operator {op_name} specified wrong suffix in parallel yaml.")
1227 return getattr(self, handler_name)(op_call, *args, **kwargs)
1229 @staticmethod
1230 def _dispatch_new(func, distribute_op, packed_call, result) -> Tensor:
1231 """New dispatch flow using preprocess result.
1233 Args:
1234 func: Original function.
1235 distribute_op: Distributed operation instance.
1236 packed_call: (prim, op_name_str) tuple for aclop kernel invocation,
1237 or None for regular ops.
1238 result: Preprocessed result (local_args, local_kwargs, cache_values).
1240 Returns:
1241 Tensor: Dispatched result as DTensor.
1242 """
1243 local_args, local_kwargs, cache_values = result
1244 cache_key = LayoutCacheKey.from_cache_values(cache_values)
1245 func_name = platform.get_op_name(func)
1247 infer_result, op_impl = OpDispatcher._lookup_or_infer_layout_new(
1248 func, func_name, cache_key, cache_values, distribute_op
1249 )
1251 op_impl = func if op_impl is None else op_impl
1252 py_output = OpDispatcher._call_op_impl(op_impl, packed_call, local_args, local_kwargs)
1253 return distribute_op.wrap_output(py_output, infer_result[0])
1255 @staticmethod
1256 def _lookup_or_infer_layout_new(func, func_name, cache_key, cache_values, distribute_op):
1257 """Look up cached layout or compute via distributed op (new three-phase API).
1259 Returns:
1260 (infer_result, op_impl)
1261 """
1262 cache_manager = LayoutCacheManager.get_instance()
1263 layout_cache = cache_manager.get_layout_cache()
1264 if func_name not in layout_cache:
1265 layout_cache[func_name] = {}
1266 op_layout_cache = layout_cache[func_name]
1267 if cache_key in op_layout_cache:
1268 return op_layout_cache[cache_key]
1269 infer_result = distribute_op.infer_layout(cache_values)
1270 op_impl = distribute_op.get_expand_impl(func, infer_result, cache_values)
1271 op_layout_cache[cache_key] = (infer_result, op_impl)
1272 return infer_result, op_impl
1274 def dispatch(self, op_call: callable, args: tuple, kwargs: dict) -> object:
1275 """Route an op call through the appropriate DTensor dispatch path.
1277 Args:
1278 op_call: The raw operator callable.
1279 args: Positional arguments for op_call.
1280 kwargs: Keyword arguments for op_call.
1282 Returns:
1283 Result of the dispatched op call.
1284 """
1285 op_name = platform.get_op_name(op_call)
1287 if self._should_bypass_dispatch(op_name):
1288 result = op_call(*self._unwrap_args(args), **self._unwrap_kwargs(kwargs))
1289 if op_name in self._INPLACE_BYPASS_OPS and args and isinstance(args[0], DTensor):
1290 return args[0]
1291 return result
1293 if op_name in self._random_ops or op_name in self._random_ms_ops:
1294 return self._dispatch_random_op(op_name, op_call, args, kwargs)
1296 self._check_decomposed_ce_op_in_loss_parallel(op_name, args, kwargs)
1298 if self._should_dispatch_loss_parallel(op_name):
1299 return self._dispatch_loss_parallel(op_call, args, kwargs)
1301 if op_name not in self.layout_infer_ops and get_distributed_op(op_name) is not None:
1302 self.layout_infer_ops[op_name] = {}
1304 return self._dispatch_layout_infer(op_name, op_call, args, kwargs)
1306_OP_DISPATCHER = OpDispatcher()