r"""Wraps a callable, and provides accelerated evaluation compiled by xla. Currently it is an experimental feature. Refer to :class:`~.jit.tracing.trace` for more information. Examples: .. code-block:: python import numpy as np from basecls.mod
| 67 | |
| 68 | |
| 69 | class xla_trace(trace): |
| 70 | r"""Wraps a callable, and provides accelerated evaluation compiled by xla. |
| 71 | Currently it is an experimental feature. |
| 72 | Refer to :class:`~.jit.tracing.trace` for more information. |
| 73 | |
| 74 | |
| 75 | Examples: |
| 76 | |
| 77 | .. code-block:: python |
| 78 | |
| 79 | import numpy as np |
| 80 | from basecls.models.resnet import resnet18 |
| 81 | from megengine.autodiff.grad_manager import GradManager |
| 82 | from megengine.jit import xla_trace |
| 83 | from megengine.optimizer import Adam |
| 84 | |
| 85 | model = resnet18() |
| 86 | gm = GradManager() |
| 87 | opt = Adam(model.parameters(), lr=1e-4) |
| 88 | gm.attach(model.parameters()) |
| 89 | |
| 90 | # Only tensors in wrapped func args/kwargs will be treated as graph inputs, |
| 91 | # and other tensors will be captured as const value. |
| 92 | # Module, optimizer, and train data/label should be arguments of the wrapped function. |
| 93 | @xla_trace(capture_as_const=True) |
| 94 | def train_step(model, opt, data, label): |
| 95 | with gm: |
| 96 | pred = model(data) |
| 97 | loss = F.loss.cross_entropy(pred, label) |
| 98 | gm.backward(loss) |
| 99 | opt.step().clear_grad() |
| 100 | return loss |
| 101 | |
| 102 | """ |
| 103 | |
| 104 | third_party_backend = True |
| 105 | |
| 106 | def __init__(self, function, *, without_host=True, symbolic_shape=False, **kwargs): |
| 107 | assert without_host, "xla trace only support without host mode" |
| 108 | assert not symbolic_shape, "xla doesn't support dynamic shape currently" |
| 109 | |
| 110 | set_external_convert_hook(apply_external_convert_hook) |
| 111 | |
| 112 | set_py_external_type(ArrayImpl) |
| 113 | set_external_convert() |
| 114 | |
| 115 | # everytime we construct a xla_trace object, which means a compilation will |
| 116 | # happen soon, so we increase `_expect_xlacompile_cnt` by one. |
| 117 | global _expect_xlacompile_cnt |
| 118 | _expect_xlacompile_cnt = _expect_xlacompile_cnt + 1 |
| 119 | |
| 120 | super().__init__( |
| 121 | function, without_host=without_host, symbolic_shape=symbolic_shape, **kwargs |
| 122 | ) |
| 123 | |
| 124 | def setup_env(self): |
| 125 | self.orig_use_xla = set_use_xla_backend(True) |
| 126 |