|
1 | 1 | import torch |
2 | | -from typing import Dict, Callable, Union, Tuple, Optional |
3 | | -from collections import defaultdict |
4 | | -import functools |
| 2 | +from typing import Union, Tuple |
5 | 3 | from dataclasses import dataclass |
6 | | -from torchao.utils import TORCH_VERSION_AT_LEAST_2_5 |
7 | | - |
8 | | -""" |
9 | | -Helper function for implementing aten op or torch function dispatch |
10 | | -and dispatching to these implementations. |
11 | | -""" |
12 | | -def _implements(cls, aten_ops_or_torch_fns): |
13 | | - """Use this decorator to implement a function for an aten ops in __torch_dispatch__ |
14 | | - (if user passed in a list of ops) |
15 | | - or torch function in __torch_function__ (if user passed in a single object) |
16 | | -
|
17 | | - class MyTensor(torch.Tensor): |
18 | | - ... |
19 | | - implements = classmethod(_implements) |
20 | | -
|
21 | | - implements = MyTensor.implements |
22 | | -
|
23 | | - @implements(torch.nn.functional.linear): |
24 | | - def _(func, types, args, kwargs): |
25 | | - ... |
26 | | -
|
27 | | - """ |
28 | | - if not hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE"): |
29 | | - cls._ATEN_OP_OR_TORCH_FN_TABLE = {} |
30 | | - |
31 | | - if not isinstance(aten_ops_or_torch_fns, (list, tuple)): |
32 | | - aten_ops_or_torch_fns = [aten_ops_or_torch_fns] |
33 | | - def decorator(func): |
34 | | - for op in aten_ops_or_torch_fns: |
35 | | - @functools.wraps(op) |
36 | | - def wrapper(f, types, args, kwargs): |
37 | | - return func(f, types, args, kwargs) |
38 | | - |
39 | | - cls._ATEN_OP_OR_TORCH_FN_TABLE[op] = wrapper |
40 | | - return func |
41 | | - return decorator |
42 | | - |
43 | | -def _dispatch__torch_function__(cls, func, types, args=(), kwargs=None): |
44 | | - """Use this util function for a common `__torch_function__` implementation |
45 | | - that dispatches to ops/functions registered with `_implements` |
46 | | -
|
47 | | - class MyTensor(torch.Tensor): |
48 | | - ... |
49 | | - __torch_function__ = classmethod(_dispatch__torch_function__) |
50 | | - """ |
51 | | - kwargs = {} if kwargs is None else kwargs |
52 | | - if hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE") and \ |
53 | | - func in cls._ATEN_OP_OR_TORCH_FN_TABLE: |
54 | | - return cls._ATEN_OP_OR_TORCH_FN_TABLE[func](func, types, args, kwargs) |
55 | | - |
56 | | - with torch._C.DisableTorchFunctionSubclass(): |
57 | | - return func(*args, **kwargs) |
58 | | - |
59 | | -def _dispatch__torch_dispatch__(cls, func, types, args, kwargs): |
60 | | - """Use this util function for a common `__torch_dispatch__` implementation |
61 | | - that dispatches to ops/functions registered with `_implements` |
62 | | -
|
63 | | - class MyTensor(torch.Tensor): |
64 | | - ... |
65 | | - __torch_dispatch__ = classmethod(_dispatch__torch_dispatch__) |
66 | | - """ |
67 | | - if hasattr(cls, "_ATEN_OP_OR_TORCH_FN_TABLE") and \ |
68 | | - func in cls._ATEN_OP_OR_TORCH_FN_TABLE: |
69 | | - return cls._ATEN_OP_OR_TORCH_FN_TABLE[func](func, types, args, kwargs) |
70 | | - |
71 | | - raise NotImplementedError(f"{cls.__name__} dispatch: attempting to run unimplemented operator/function: {func}") |
72 | | - |
73 | 4 |
|
74 | 5 | """ |
75 | 6 | Base class for different LayoutType, should not be instantiated directly |
@@ -101,52 +32,6 @@ def extra_repr(self) -> str: |
101 | 32 | class PlainLayoutType(LayoutType): |
102 | 33 | pass |
103 | 34 |
|
104 | | -""" |
105 | | -layout tensor constructor registration for different tensor subclassesa |
106 | | -
|
107 | | -first key is a tensor subclass type like AffineQuantizedTensor |
108 | | -second key is an extended layout string, like tensor_core_tiled |
109 | | -value is a constructor for the LayoutTensor class, e.g. TensorCoreTiledAQTLayout.from_plain |
110 | | -""" |
111 | | -_LAYOUT_CONSTRUCTOR_TABLE: Dict[Callable, Dict[type(LayoutType), Callable]] = defaultdict(dict) |
112 | | - |
113 | | -def _register_layout_cls(cls: Callable, layout_type_class: type(LayoutType)): |
114 | | - """Helper function for layout registrations, this is used to implement |
115 | | - register_layout_cls decorator for each tensor subclass, see aqt.py for example usage |
116 | | -
|
117 | | - Args: |
118 | | - cls: Tensor subclass type |
119 | | - layout_type_class: the class type of subclass of `LayoutType`, e.g. `PlainLayoutType` |
120 | | -
|
121 | | - Returns: |
122 | | - a decorator that registers the layout tensor constructor in the table |
123 | | - """ |
124 | | - def decorator(layout_cls): |
125 | | - _LAYOUT_CONSTRUCTOR_TABLE[cls][layout_type_class] = layout_cls.from_plain |
126 | | - if TORCH_VERSION_AT_LEAST_2_5: |
127 | | - # Allow serialization to work for models uses this layout tensor subclass |
128 | | - torch.serialization.add_safe_globals([layout_type_class, layout_cls]) |
129 | | - return layout_cls |
130 | | - return decorator |
131 | | - |
132 | | -def _get_layout_tensor_constructor(cls: Callable, layout_type_class: type(LayoutType)) -> Callable: |
133 | | - """Get Layout class constructor (LayoutClass.from_plain) for `cls` based on `layout_type_class` |
134 | | - `layout_type_class` means the class type of subclass of `LayoutType`, e.g. `PlainLayoutType` |
135 | | -
|
136 | | - Args: |
137 | | - cls: Tensor subclass type |
138 | | - layout_type_class: the class type of subclass of `LayoutType`, e.g. `PlainLayoutType` |
139 | | -
|
140 | | - Returns: |
141 | | - layout tensor subclass constructor for the layout_type_class |
142 | | - """ |
143 | | - if cls not in _LAYOUT_CONSTRUCTOR_TABLE: |
144 | | - raise ValueError(f"no registered layout class constructor for: {cls}") |
145 | | - if layout_type_class not in _LAYOUT_CONSTRUCTOR_TABLE[cls]: |
146 | | - raise ValueError(f"layout_name: {layout_type_class} is not supported yet for {cls}") |
147 | | - |
148 | | - return _LAYOUT_CONSTRUCTOR_TABLE[cls][layout_type_class] |
149 | | - |
150 | 35 | def is_device(target_device_str: str, device: Union[str, torch.device]): |
151 | 36 | return torch.device(device).type == target_device_str |
152 | 37 |
|
|
0 commit comments