Repository navigation
Expand file tree
/
Copy pathgraph.py
More file actions
1874 lines (1593 loc) · 74.7 KB
/
Copy pathgraph.py
File metadata and controls
1874 lines (1593 loc) · 74.7 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# ===----------------------------------------------------------------------=== #
# Copyright (c) 2026, Modular Inc. All rights reserved.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions:
# https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===----------------------------------------------------------------------=== #
"""Core graph primitives."""
from __future__ import annotations
import contextlib
import functools
import inspect
import io
import itertools
import sys
import traceback
from collections import OrderedDict
from collections.abc import Callable, Generator, Iterable, Mapping, Sequence
from contextvars import ContextVar
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Any, TypeAlias, TypeGuard, TypeVar, cast
from max import mlir, tree
from max._core import Attribute as _Attribute
from max._core import Block, OpBuilder, Operation
from max._core import Type as _Type
from max._core import Value as _Value
from max._core import graph as _graph
from max._core.dialects import builtin, kgen
from max._core.dialects import kgen as _kgen
from max._core.dialects import mo as _mo
from max._core.dialects.m import DeviceInfoAttr as _DeviceInfoAttr
from max._core.driver import CPU, Accelerator, Device, accelerator_count
from max._core.engine import InferenceSession as _InferenceSession
from max._mlir_context import (
default_mlir_context,
ensure_default_mlir_context,
in_default_mlir_context,
)
from max.mlir.dialects import mo
from mojo.paths import (
_build_mojo_source_package,
is_mojo_binary_package_path,
is_mojo_source_package_path,
)
from ._api_boundary import api_boundary, callers_of_api
from .type import (
BufferType,
DeviceRef,
SymbolicDim,
TensorType,
Type,
_ChainType,
)
from .value import BufferValue, TensorValue, TensorValueLike, Value, _ChainValue
GraphOutputs: TypeAlias = (
Value[Any]
| TensorValueLike
| Sequence["GraphOutputs"]
| Mapping[Any, "GraphOutputs"]
)
"""What a graph outputs: values, or a tree of them whose leaves are the
outputs in order. Sequences and mappings rather than :data:`max.tree.Tree`'s
lists and dicts, so a ``list[TensorValue]`` or a typed group is accepted.
Mapping keys only name the outputs, so any key type is."""
from .weight import Weight
CURRENT_GRAPH: ContextVar[Graph] = ContextVar("CURRENT_GRAPH")
# Stack of active Graph.profile_scope() scopes, outermost first. Each entry is
# a (name, color) tuple; color is None when unset. A tuple (not a list) so
# entering/exiting a scope is a ContextVar set()/reset() rather than a
# mutation shared across contexts.
# NOTE: this is a single module-level ContextVar, not scoped per Graph.
# Building a nested Graph while a sibling graph's profile_scope is active will
# inherit that scope's label; this is not currently supported and should be
# avoided.
_CURRENT_PROFILE_SCOPES: ContextVar[tuple[tuple[str, str | None], ...]] = (
ContextVar("_CURRENT_PROFILE_SCOPES", default=())
)
_KERNEL_LIBRARY_PATHS_ATTR_NAME = "_kernel_library_paths"
_DEVICE_INFO_MAPPING_ATTR_NAME = "mo.device_info_mapping"
T = TypeVar("T")
class ProfileScopeColor(str, Enum):
"""Named color for the optional in-region ``Graph.profile_scope`` range.
These names mirror the palette accepted by the GPU tracing backend.
"""
MODULAR_PURPLE = "modular_purple"
"""The Modular brand purple."""
BLUE = "blue"
"""The color blue."""
GREEN = "green"
"""The color green."""
ORANGE = "orange"
"""The color orange."""
PURPLE = "purple"
"""The color purple."""
RED = "red"
"""The color red."""
WHITE = "white"
"""The color white."""
YELLOW = "yellow"
"""The color yellow."""
def _is_chain_value(value: _Value[Any]) -> TypeGuard[_Value[_mo.ChainType]]:
return isinstance(value.type, _mo.ChainType)
class _DeviceChainMap(OrderedDict[DeviceRef, _ChainValue]):
"""Dictionary that lazily registers per-device chains on access."""
def __init__(self, graph: Graph) -> None:
super().__init__()
self._graph = graph
@staticmethod
def _sort_key(device: DeviceRef) -> tuple[str, int]:
return (device.device_type.value, device.id)
def __getitem__(self, device: DeviceRef) -> _ChainValue:
if device not in self:
self[device] = (
self._graph._add_chain_block_arg()
if self._graph._has_chain_input
# prevent block arguments from being added when we're inside a
# region of a control flow op (mo.if, mo.while). For simplicity
# currently control flow ops merge new device chains with the
# global chain instead of shuttling device chains from the
# parent scope via block arguments. Any device chains that
# exist prior to the control flow op will be used as expected.
and self._graph._graph_body == self._graph._current_block
# Seed new device chains from the host (``DeviceRef.CPU()``)
# chain, which is pre-populated at graph construction.
else super().__getitem__(DeviceRef.CPU())
)
return super().__getitem__(device)
def __setitem__(self, device: DeviceRef, chain: _ChainValue) -> None:
assert isinstance(chain, _ChainValue)
super().__setitem__(device, chain)
def __iter__(self):
return iter(sorted(super().keys(), key=self._sort_key))
def _value(self, device: DeviceRef) -> _ChainValue:
return super().__getitem__(device)
def ordered_values(self) -> Generator[_ChainValue]:
for device in self:
yield self._value(device)
def pack(self, values: Iterable[Any] = ()) -> list[Any]:
"""Append the current chain state to ``values``.
Returns ``[*values, *self.ordered_values()]`` — useful for bundling
user values with the chain state for control-flow op operands /
block-arg signatures / terminator operands.
"""
return [*values, *self.ordered_values()]
def unpack(self, values: Sequence[Any]) -> list[Any]:
"""Take trailing chain values off ``values`` and return user values.
Inverse of :meth:`pack`: assumes the last ``len(self)`` entries of
``values`` are chain values in the same order as iteration over
``self``, assigns them back to ``self[device]``, and returns the
rest. Side-effecting on ``self``.
"""
chain_count = len(self)
user_values = list(values[:-chain_count] if chain_count else values)
chain_values = values[-chain_count:] if chain_count else ()
for device, chain in zip(self, chain_values, strict=True):
assert isinstance(chain, _ChainValue)
self[device] = chain
return user_values
def merge_for(self, devices: Iterable[DeviceRef]) -> _ChainValue:
"""Merge the host chain and per-device chains for ``devices`` into one.
Multi-device collective ops need a single combined input chain that
depends on the host orchestration chain and on each participating
device's compute chain. This emits the ``mo.chain.create`` that
combines them, updates the host chain to the merged result, and
returns it.
Duplicates are removed (for example, if ``DeviceRef.CPU()`` appears in
``devices``, the host chain isn't listed twice).
"""
keys = list(dict.fromkeys((DeviceRef.CPU(), *devices)))
chains = [self[device] for device in keys]
unique_chains = list(dict.fromkeys(chains))
if len(unique_chains) == 1:
[chain] = unique_chains
assert chain == self[DeviceRef.CPU()]
return chain
merged = self._graph._add_op_generated(
_mo.ChainCreateOp,
result=_mo.ChainType(),
inputs=unique_chains,
attach_profile_scopes=False,
)[0]
assert isinstance(merged, _ChainValue)
self[DeviceRef.CPU()] = merged
return merged
def copy(self) -> _DeviceChainMap:
result = _DeviceChainMap(self._graph)
for device in self:
result[device] = self._value(device)
return result
def __repr__(self) -> str:
items = ", ".join(f"{device}: {self._value(device)}" for device in self)
return f"_DeviceChainMap({{{items}}})"
class KernelLibrary:
"""Manages custom kernel libraries and operations for a graph.
A kernel library provides access to custom operations and kernels that can
be loaded from various sources including Mojo pre-compiled packages (``.mojoc``)
and Mojo source directories. The library handles verification and registration
of custom operations within the MLIR context.
"""
_analysis: _graph.Analysis
_analysis_cache: dict[frozenset[str], _graph.Analysis] = {}
def __init__(self, paths: Iterable[Path] = ()) -> None:
context = default_mlir_context()
_graph._init_and_register_max_context(context)
paths_list = list(paths)
self._analysis = _graph.Analysis(context, paths_list)
def library_paths(self) -> list[Path]:
"""Returns the list of kernel library paths.
Returns:
A list of :class:`pathlib.Path` objects representing the currently loaded
kernel library paths.
"""
return self._analysis.library_paths
def add_path(self, path: Path) -> None:
"""Adds a kernel library path to the analysis.
Args:
path: The :class:`pathlib.Path` to the kernel library to be added to the
current analysis.
"""
self._analysis.add_path(path)
def load_paths(self, custom_extensions: Iterable[Path]) -> None:
"""Loads custom operations from provided library paths.
Performs additional "smart" library loading logic for custom operation
libraries in additional formats. The loading logic supports the
following formats:
- Compiled Mojo binary packages with ``.mojoc`` extension
- Mojo source directory with custom operations
The loaded libraries are added to the current kernel library.
Args:
custom_extensions: The file paths to the custom operation libraries.
"""
for ext_path in custom_extensions:
if is_mojo_binary_package_path(ext_path):
self.add_path(ext_path)
elif is_mojo_source_package_path(ext_path):
# Builds the source directory into a precompiled Mojo file.
self.add_path(_build_mojo_source_package(ext_path))
else:
raise ValueError(
"Path provided as custom extension to Graph must be a "
+ f"Mojo source or binary package: {ext_path}"
)
def __getitem__(self, kernel: str):
if kernel not in self._analysis.symbol_names:
raise KeyError(kernel)
return self._analysis.kernel(kernel)
def __contains__(self, kernel: str) -> bool:
return kernel in self._analysis.symbol_names
def __iter__(self):
yield from sorted(self._analysis.symbol_names)
def verify_custom_op(self, custom_op: mlir.Operation) -> None:
"""Verifies that a custom operation is valid within the current context.
Args:
custom_op: The :obj:`mlir.Operation` to be verified against the
current kernel library analysis.
"""
paths = self._analysis.library_paths
if paths:
cache_key = frozenset(str(p) for p in paths)
cached = KernelLibrary._analysis_cache.get(cache_key)
if cached is not None:
self._analysis = cached
self._analysis.verify_custom_op(custom_op)
return
self._analysis.verify_custom_op(custom_op)
KernelLibrary._analysis_cache[cache_key] = self._analysis
else:
self._analysis.verify_custom_op(custom_op)
def has_shape_function(self, kernel: str) -> bool:
"""Returns whether *kernel* registers a shape function.
A kernel that registers no shape function has no way to compute its
output shape at run time, so the graph compiler rejects a
data-dependent output dimension declared for it.
Args:
kernel: The registered name of the kernel to check.
Returns:
``True`` if a shape function is registered for *kernel*.
Raises:
KeyError: If no kernel named *kernel* is in the library.
"""
if kernel not in self:
raise KeyError(f"no kernel named {kernel!r} in the kernel library")
return self._analysis.has_shape_function(kernel)
_default_custom_extensions: tuple[Path, ...] = ()
def default_custom_extensions() -> tuple[Path, ...]:
"""Returns the custom-extension paths implicitly loaded by every new graph.
A backend whose ops need a kernel-overlay library registers it here so ops
resolve to the overlays even on graph paths that don't thread
``custom_extensions`` explicitly — notably the experimental
eager-realization ``Graph("main")``. Empty by default, so there is no
effect unless a backend registers a library.
"""
return _default_custom_extensions
def _resolved_custom_extensions(
custom_extensions: Iterable[Path],
) -> list[Path]:
"""Unions *custom_extensions* with the process-global overlay.
Explicit paths keep their order and come first; an overlay path already
given explicitly is not repeated. This is the one place that ordering
lives, so anything that has to agree with what a :class:`Graph` links
against (a shape-function probe, a compiled-binding cache key) resolves
through here rather than combining the two lists its own way.
"""
extensions = list(custom_extensions)
extensions.extend(
path for path in default_custom_extensions() if path not in extensions
)
return extensions
@contextlib.contextmanager
def default_custom_extensions_scope(*paths: Path) -> Generator[None]:
"""Adds *paths* to :func:`default_custom_extensions` for the block's duration.
Paths already registered are not duplicated. The previous defaults are
restored on exit.
"""
global _default_custom_extensions
previous = _default_custom_extensions
merged = list(previous)
merged.extend(path for path in paths if path not in merged)
_default_custom_extensions = tuple(merged)
try:
yield
finally:
_default_custom_extensions = previous
class DevicePlacementPolicy(Enum):
"""Controls behavior when an op implicitly transfers a tensor to CPU.
Ops that only have CPU kernels must transfer non-CPU tensors before
executing. This policy controls how that situation is reported:
- ``Ignore``: transfer silently, no message.
- ``Warn`` (default): emit a ``UserWarning`` naming the op and the
tracking ticket for GPU support.
- ``Error``: raise ``ValueError``, making the implicit transfer a hard
build-time failure.
Pass via ``Graph(..., strict_device_placement=DevicePlacementPolicy.Error)``.
"""
Ignore = "ignore"
"""Transfer silently, without any message."""
Warn = "warn"
"""Emit a ``UserWarning`` naming the op; the default policy."""
Error = "error"
"""Raise ``ValueError``, making the implicit transfer a hard build-time
failure."""
# From https://stackoverflow.com/a/76301341
class _classproperty:
def __init__(self, func) -> None: # noqa: ANN001
self.fget = func
def __get__(self, instance, owner): # noqa: ANN001
return self.fget(owner)
@dataclass(frozen=True)
class _GraphWeight:
weight: Weight
value: TensorValue
def _location(
ignore_frames: int = 1, attach_profile_scopes: bool = True
) -> mlir.Location:
"""Creates an MLIR Location with the current Python call stack."""
if not mlir.Context.current:
raise RuntimeError("Can't create location: No MLIR context active")
# Read the live config value rather than caching at import time: the
# flag can be flipped after import via Graph.debug.source_tracebacks,
# InferenceSession.debug.sensible_mode, or MODULAR_DEBUG, and each
# graph op must honor the setting at the moment it is created.
if not _InferenceSession.debug.source_tracebacks:
location: mlir.Location = mlir.Location.unknown()
else:
# Extract the stack into summaries
# - Avoids reference cycles
# - Doesn't keep references to closures
# Below the call into the graph API, the frames are MAX's
# implementation of the op. An op built through an internal path has
# no boundary on its stack and keeps every frame.
callers = (
callers_of_api(sys._getframe(1))
if _InferenceSession.debug.source_tracebacks_hide_internal
else None
)
tb: list[traceback.FrameSummary]
if callers is not None:
tb = list(
traceback.StackSummary.extract(
(frame, frame.f_lineno) for frame in callers
)
)
else:
# Always remove at least _location
tb = traceback.extract_stack()[: -(ignore_frames + 1)]
location = (
_graph.frame_loc(mlir.Context.current, tb)
if tb
else mlir.Location.unknown()
)
if attach_profile_scopes:
# Wrap outermost scope first, so the innermost scope ends up as the
# outermost ProfileScopeLocationAttr and the location assembly reads
# in nesting order. This runs regardless of source-traceback capture:
# profile_scope labels are a distinct, always-on mechanism riding the
# same Location, not a feature of the debug traceback.
for name, color in _CURRENT_PROFILE_SCOPES.get():
location = _graph.profile_scope_location(
mlir.Context.current, name, location, color
)
return location
def _to_mlir(o: Any) -> Any:
# Convert args from instances of Python graph-api Value() to mlir.Value
if hasattr(o, "to_mlir"):
return o.to_mlir()
elif isinstance(o, list | tuple):
return type(o)(_to_mlir(ov) for ov in o)
elif isinstance(o, dict):
return {k: _to_mlir(v) for k, v in o.items()}
return o
# MLIR assembly text must be read and written through the ``max.mlir`` bindings
# rather than ``max._core``. ``_core`` static-links its own LLVM, so its textual
# float parser and printer compare ``llvm::APFloat`` semantics against their own
# copy of the ``APFloatBase::sem*`` statics and never match the ones the
# attribute was built against. That silently turns f64 constants into denormal
# garbage and aborts on any f32 that MLIR spells as a hex literal. ``max.mlir``
# resolves MLIR from ``libmax``, the copy that owns the context, so a round-trip
# through it agrees with itself. See GEX-4052.
def _asm_via_cmlir(op: Operation) -> str:
"""Serializes an operation to MLIR assembly text.
The cmlir wrapper is a non-owning view onto ``op``, so whatever owns ``op``
keeps owning it.
"""
return str(mlir.Operation._CAPICreate(op._CAPIPtr))
def _parse_module_via_cmlir(
asm: str, context: mlir.Context
) -> builtin.ModuleOp:
"""Parses MLIR assembly text into a typed ``builtin.ModuleOp``.
Parses in ``max.mlir``, then moves the module into ``max._core`` as
bytecode, which is a lossless transport between the two. The result is
detached, matching what ``_core.parse_module`` returns.
"""
parsed = mlir.Module.parse(asm, context)
bytecode = io.BytesIO()
parsed.operation.write_bytecode(bytecode)
return Operation.from_bytecode(bytecode.getvalue(), context)
def _set_output_param_decls(op: Operation, params: dict[str, None]) -> None:
# Interfaces don't yet support isinstance checks, so this is a cheap proxy.
# - nanobind doesn't allow custom metaclasses, but __instancecheck__
# must be defined on a metaclass
# - Interfaces are protocols even though we know when they are explicitly
# implemented so that attrs/types/ops may implement them without declaring
# it in the stub files
# - it's not trivial to define our own `isa`-like check. Such a check needs
# to name the template parameter for `mlir::isa<T>` explicitly in C++, but
# if `isa` is defined as a staticmethod it interferes with protocol type
# checking.
if not hasattr(op, "output_param_decls"):
return
op: Operation & _mo.ParamDeclarationInterface # type: ignore
# Add symbolic dims of tensor results to the list of graph params and
# declared output params of the op
# Use a dict as an ordered set for new param decls. Maps keys to None.
result_parameters = dict.fromkeys(
itertools.chain.from_iterable(
value.type.parameters
for result in op.results
if isinstance(
value := Value.from_mlir(result), TensorValue | BufferValue
)
)
)
names = [parameter.name for parameter in result_parameters]
# Track newly declared parameters. Not `names - params.keys()`: set order
# is per-process hash order and lands in the MEF cache key (MXF-584).
if new_params := dict.fromkeys(n for n in names if n not in params):
params.update(new_params)
si64 = kgen.SIMDType(1, kgen._KGENDType.get_int(64, True))
# We can't overload the setter yet, so the interface annotation is wrong
# TODO(MAXPLAT-306): See https://github.com/wjakob/nanobind/discussions/1063
op.output_param_decls = kgen.ParamDeclArrayAttr(
[kgen.ParamDeclAttr(name, si64) for name in new_params]
)
class Module:
"""A container for one or more :class:`Graph` instances compiled together.
A ``Module`` holds an MLIR module that may contain multiple ``mo.graph``
ops — for example a vision encoder and a language model that should be
handed to :meth:`max.engine.InferenceSession.load_all` as a single
compilation unit. Constructing several :class:`Graph` instances with the
same ``module=`` argument adds each graph as a top-level op in this
module.
For example, build two graphs into a shared module, load them in one
compilation unit, then execute both:
.. code-block:: python
import numpy as np
from max.driver import Accelerator, CPU, accelerator_count
from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, Module, ops
device = Accelerator() if accelerator_count() > 0 else CPU()
device_ref = DeviceRef.from_device(device)
module = Module()
with Graph("encoder", input_types=[], module=module) as encoder:
encoder.output(
ops.constant(
np.array([1.0, 2.0, 3.0], dtype=np.float32),
DType.float32,
device=device_ref,
)
)
with Graph("decoder", input_types=[], module=module) as decoder:
decoder.output(
ops.constant(
np.array([4.0, 5.0, 6.0], dtype=np.float32),
DType.float32,
device=device_ref,
)
)
models = InferenceSession(devices=[device]).load_all(module)
(enc_out,) = models[encoder.name].execute()
(dec_out,) = models[decoder.name].execute()
.. invisible-code-block: python
np.testing.assert_allclose(enc_out.to_numpy(), [1.0, 2.0, 3.0])
np.testing.assert_allclose(dec_out.to_numpy(), [4.0, 5.0, 6.0])
The wrapped MLIR module is exposed as :attr:`mlir_module` for code that
must reach the underlying representation (graph compiler internals,
serializers and so on); routine users should not need it.
"""
mlir_module: builtin.ModuleOp
"""The wrapped MLIR module, exposed through the typed nanobind dialect
bindings."""
def __init__(self, mlir_module: builtin.ModuleOp | None = None) -> None:
"""Create a new, empty module, or wrap an existing :class:`builtin.ModuleOp`.
Args:
mlir_module: An existing typed :class:`builtin.ModuleOp` to
wrap. When ``None`` (the default), a new empty module is
created. The wrap-existing form is intended for code that
already holds a typed ``ModuleOp`` (for example
:class:`Graph`'s subgraph plumbing or the
:attr:`Graph.module` property) and needs to expose it
through the public :class:`Module` surface; routine users
should leave this argument unset.
"""
if mlir_module is None:
mlir_module = builtin.ModuleOp(location=_location())
self.mlir_module = mlir_module
def top_level_graph_names(self) -> list[str]:
"""Return the name of every top-level (non-subgraph) graph in this module.
Walks the wrapped module's body and skips any op whose
``is_subgraph`` is set (subgraphs are inlined callees, not loaded
as standalone models). Non-graph ops are skipped rather than
raising. The resulting order matches MEF model order, which is the
order :func:`max.engine.InferenceSession.load_all` returns models
in.
Returns:
The names of the top-level graphs, in MEF order. The list is
empty for a freshly constructed :class:`Module`.
"""
names: list[str] = []
for op in self.mlir_module.body:
if isinstance(op, _mo.GraphOp) and not op.is_subgraph:
names.append(op.sym_name)
return names
def _to_mlir_str(self, *, source_locations: bool = False) -> str:
"""Serializes this module to MLIR assembly text.
Internal helper used by graph-dump tooling.
Args:
source_locations: When ``True``, annotates each op with the Python
call stack it was built from. This requires source-traceback
capture to have been enabled during graph construction (see
:attr:`max.graph.Graph.debug`); without it the ops carry no
Python frames and the annotations are empty. The wrapped module
is left unchanged either way.
Returns:
The module's MLIR assembly text.
"""
if not source_locations:
return _asm_via_cmlir(self.mlir_module)
# The materialized locations come back as bytecode; reparsing through
# max.mlir is what puts the printing in the right LLVM image.
with_locations = mlir.Module.parse(
_graph.source_locations_bytecode(self.mlir_module),
self.mlir_module.context,
)
return with_locations.operation.get_asm(enable_debug_info=True)
class GraphDebugConfig:
"""Narrow view of :class:`max.engine.DebugConfig` exposed through :attr:`Graph.debug`.
The attribute :attr:`source_tracebacks` lives on ``Graph.debug`` because it is
consumed during graph construction, before an ``InferenceSession`` exists.
All other debug options are available on ``InferenceSession.debug`` and
share the same global state.
"""
@property
def source_tracebacks(self) -> bool:
"""See :attr:`max.engine.DebugConfig.source_tracebacks`."""
return _InferenceSession.debug.source_tracebacks
@source_tracebacks.setter
def source_tracebacks(self, value: bool) -> None:
_InferenceSession.debug.source_tracebacks = value
@property
def source_tracebacks_hide_internal(self) -> bool:
"""See :attr:`max.engine.DebugConfig.source_tracebacks_hide_internal`."""
return _InferenceSession.debug.source_tracebacks_hide_internal
@source_tracebacks_hide_internal.setter
def source_tracebacks_hide_internal(self, value: bool) -> None:
_InferenceSession.debug.source_tracebacks_hide_internal = value
class Graph:
"""Represents a single MAX graph.
A :class:`Graph` defines a model's computation. You build a graph by
composing operations that describe how input tensors are transformed into
outputs. Unlike imperative code that executes operations, a :class:`Graph`
captures the data flow between operations, which allows MAX to optimize and
parallelize execution at compile time. Operations run on the compiled object.
The following code examples show two different strategies for constructing
graphs.
**Use the context manager:** Use :class:`Graph` as a context manager to
define the active graph. Inside the ``with`` block, retrieve inputs from
:attr:`inputs`, call ops to build nodes, and set the graph output with
:meth:`output()`. Ops called inside the block find the active graph
automatically. Ops called outside the block fail because there is no active
graph.
.. code-block:: python
from max.dtype import DType
from max.graph import DeviceRef, Graph, TensorType, Weight
W = Weight("W", DType.float32, [3, 2], DeviceRef.CPU())
b = Weight("b", DType.float32, [2], DeviceRef.CPU())
with Graph(
"linear_relu",
input_types=[TensorType(DType.float32, ["batch", 3], device=DeviceRef.CPU())],
) as graph:
x = graph.inputs[0].tensor
y = x @ W + b
graph.output(y)
**Use the graph constructor:** Pass a callable as the ``forward`` argument.
The graph automatically passes the input :class:`TensorValue` to the
callable and records the return value as the graph output. Under the hood,
this still opens and closes a graph context.
.. code-block:: python
from max.dtype import DType
from max.graph import DeviceRef, Graph, TensorType, TensorValue, Weight, ops
class Linear:
def __init__(self, in_dim: int, out_dim: int):
self.weight = Weight("W", DType.float32, [in_dim, out_dim], DeviceRef.CPU())
self.bias = Weight("b", DType.float32, [out_dim], DeviceRef.CPU())
def __call__(self, x: TensorValue) -> TensorValue:
return ops.matmul(x, self.weight) + self.bias
linear_layer = Linear(2, 2)
graph = Graph(
"linear",
linear_layer,
input_types=[TensorType(DType.float32, (2,), DeviceRef.CPU())],
)
These examples only use the :obj:`max.graph` package, but most models also
use :class:`~max.nn.Module` and other building blocks from :obj:`max.nn`.
To learn more, see `Build a model graph with Module
</develop/modules>`_.
Args:
name: A name for the graph.
forward: The sequence of graph ops for the forward pass (inference).
input_types: A sequence of :class:`~max.graph.type.Type` instances that
describe each graph input.
These are typically :class:`TensorType` instances. You can also
include :class:`BufferType` instances for mutable in-place inputs.
path: The path to a saved graph (internal use only).
custom_extensions: The extensions to load for the model. Supports paths
to ``.mojoc`` or ``.mojo`` sources with custom ops.
kernel_library: Optional pre-built kernel library to use. Defaults to
``None`` (a new library is created from ``custom_extensions`` if
needed).
module: Optional existing MLIR module (internal use only). Defaults to
``None``.
"""
debug = GraphDebugConfig()
# Use a dict rather than a set to keep params ordered.
# This is to make IR generation deterministic for model IR cache hits.
# Note that insertion order in built-in dict has been guaranteed since
# Python 3.7.
_params: dict[str, None]
_mlir_op: mlir.Operation | mlir.OpView
_graph_body: mlir.Block
_module: builtin.ModuleOp
_context_state: list[contextlib.AbstractContextManager[Graph]]
_weights: dict[str, _GraphWeight]
_current_block: mlir.Block
_should_verify_ops: bool
_has_chain_input: bool
# How the value passed to :meth:`output` nests the outputs; ``None`` for a
# graph read back from MLIR, which records no structure.
_output_structure: tree.TreeDef | None = None
# Per-device chains that ensure the correct sequence of device execution.
# ``device_chains[DeviceRef.CPU()]`` is the host orchestration chain (the
# chain advanced by host-side ops like ``debug.print``, ``mo.call``, and
# the host side of collectives); all other entries are per-device compute
# chains. New device entries are seeded from the host chain.
device_chains: _DeviceChainMap
_kernel_library: KernelLibrary
_subgraphs: dict[str, Graph] = {}
# Ensure the default MLIR context for bg-thread Graph construction.
@in_default_mlir_context
def __init__(
self,
name: str,
forward: Callable[..., GraphOutputs | None] | None = None,
input_types: Iterable[Type[Any]] = (),
path: Path | None = None,
*args,
custom_extensions: Iterable[Path] = [],
kernel_library: KernelLibrary | None = None,
module: Module | None = None,
strict_device_placement: DevicePlacementPolicy = DevicePlacementPolicy.Warn,
is_device_graph: bool = False,
**kwargs,
) -> None:
self.name = name
self.strict_device_placement = strict_device_placement
if path is not None:
self._load_mlir(path)
return
self._params = dict.fromkeys(
dim.name
for t in input_types
if isinstance(t, TensorType | BufferType)
for dim in t.shape
if isinstance(dim, SymbolicDim)
)
self._context_state = []
self._data_dependent_dims = 0
self._should_verify_ops = True
with _location() as loc:
# The top-level module op is stored as a typed nanobind
# :class:`builtin.ModuleOp`. Either reuse the one supplied via
# ``module=`` or create a fresh one.
self._module = (
module.mlir_module if module else builtin.ModuleOp(location=loc)
)
builder = OpBuilder(self._module.body.end)
op = _mo.GraphOp(
builder,
loc,
name=name,
input_types=[t.to_mlir() for t in input_types],
result_types=[],
)
si64 = kgen.SIMDType(1, kgen._KGENDType.get_int(64, True))
# TODO(MAXPLAT-306): Type annotations are wrong here
op.input_parameters = kgen.ParamDeclArrayAttr(
[kgen.ParamDeclAttr(p, si64) for p in self._params]
)
self._mlir_op = mlir.Operation._CAPICreate(op._CAPIPtr)
self._current_block = self._mlir_op.regions[0].blocks[0]
self._graph_body = self._current_block
self._populate_device_info_mapping()
self._weights = {}
self._has_chain_input = False
self.device_chains = _DeviceChainMap(self)
# Create an always-ready chain that is never advanced by the graph.
# Use it for operations that are safe to schedule without per-device
# ordering constraints (e.g., host→device transfers for staging).
self._always_ready_chain = _ChainValue(
self._add_op_generated(
_mo.ChainCreateOp,
result=_mo.ChainType(),
inputs=[],
attach_profile_scopes=False,
)[0]
)
self._update_chain(self._always_ready_chain)
if self._graph_body.arguments:
mlir_maybe_chain_value = _Value._from_cmlir(
self._graph_body.arguments[-1]
)
if _is_chain_value(mlir_maybe_chain_value):
self._has_chain_input = True
self._update_chain(
_ChainValue.from_mlir(mlir_maybe_chain_value)
)
# Resolving through `_resolved_custom_extensions` is what lets a graph
# built without `custom_extensions` still reach a backend's kernel
# overlays (see `default_custom_extensions`).
self._kernel_library = kernel_library or KernelLibrary()
self._import_kernels(_resolved_custom_extensions(custom_extensions))
self._subgraphs = {}
# If we're building a device graph, annotate the graph appropriately.
if is_device_graph:
op.is_device_graph = builtin.UnitAttr()
if forward is not None:
# If the forward method was passed stage the graph directly in the
# constructor.
with self:
result = forward(*self.inputs, *args, **kwargs)
# Forward returns None, one value, a tree such as a tuple or a
# named dict, or any other iterable of values.
if result is None:
self.output()
elif tree.is_node(result):
self.output(result)
elif isinstance(result, Iterable):
self.output(*result)
else:
self.output(result)
@functools.cached_property
def inputs(self) -> Sequence[Value[Any]]:
"""The input values of the graph.
Returns:
A sequence of :class:`~max.graph.Value` objects corresponding to
the ``input_types`` passed at construction, excluding internal
chain values.
"""
body_args = self._graph_body.arguments
chain_count = len(self.device_chains) if self._has_chain_input else 0
if body_args and chain_count:
body_args = body_args[:-chain_count]
return tuple(
Value.from_mlir(_Value._from_cmlir(arg)) for arg in body_args
)
def add_subgraph(
self,
name: str,
forward: Callable[..., GraphOutputs | None] | None = None,
input_types: Iterable[Type[Any]] = (),
path: Path | None = None,
custom_extensions: Iterable[Path] = [],
devices: Iterable[DeviceRef] = [],
is_device_graph: bool = False,
) -> Graph:
"""Creates a reusable subgraph for the current graph.
A subgraph is the graph equivalent of a function: you define a block of
ops once and call it from the parent graph as many times as you need.
Use a subgraph when a block of computation repeats, for example, a
transformer layer that appears 62 times in a model. Wrapping it in a
subgraph lets the compiler process the definition once instead of once
per repetition, which can cut compile time by 50x or more.
Trade-offs to keep in mind:
- **Memory:** Allocations inside a subgraph can't be shared with