Adding Hlo module parsing util functions.

jax_scalify/__init__.py

@@ -1,5 +1,5 @@
 # Copyright (c) 2023 Graphcore Ltd. All rights reserved.
-from . import core, lax, ops, tree
+from . import core, lax, ops, tree, utils
 from .core import ( # noqa: F401
  Pow2RoundMode,
  ScaledArray,

jax_scalify/utils/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+from .hlo import parse_hlo_module, print_hlo_module # noqa: F401

jax_scalify/utils/hlo.py

@@ -0,0 +1,107 @@
+# Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+import json
+from dataclasses import dataclass
+from typing import Any, Dict, List
+import textwrap
+
+from jax.stages import Compiled, Lowered
+
+
+@dataclass
+class HloOperationInfo:
+ """HLO module operation (raw) info.
+
+ Parse from raw `as_text` compiled HloModule.
+
+ Args:
+ cmd: Raw HLO operation (function + inputs/outputs).
+ metadata: JAX metadata (line, ...)
+ backend_config: Optional backend config dictionary.
+ """
+
+ cmd: str
+ indent: int = 0
+ metadata: str | None = None
+ backend_config: Dict[Any, Any] | None = None
+
+ def as_text(self, metadata: bool = False, backend_cfg: bool = False, indent: int = 2) -> str:
+ """Convert to raw text, with formatting issues."""
+ indent_txt = " " * (indent * self.indent)
+ line = indent_txt + self.cmd
+ if backend_cfg and self.backend_config:
+ # A bit hacky text formating of backend config!
+ backend_cfg_raw = json.dumps(self.backend_config, indent=indent)
+ backend_cfg_raw = "backend_cfg: " + backend_cfg_raw
+ backend_cfg_raw = textwrap.indent(backend_cfg_raw, indent_txt + " " * indent)
+ line += "\n" + backend_cfg_raw
+ return line
+
+
+def parse_hlo_operation_raw_line(raw_line: str) -> HloOperationInfo:
+ """Very crude and ugly parsing of an Hlo operation raw line!
+
+ Returns:
+ Parsed Hlo operation line.
+ """
+ metadata: str | None = None
+ backend_cfg = None
+
+ # Parse "metadata={...}" block.
+ metadata_prefix = ", metadata={"
+ lidx = raw_line.find(metadata_prefix)
+ if lidx >= 0:
+ ridx = raw_line[lidx:].find("}") + lidx
+ metadata = raw_line[lidx : ridx + 1]
+ raw_line = raw_line.replace(metadata, "")
+ metadata = metadata[2:]
+
+ # Parse "backend_config={...}" block.
+ backend_cfg_prefix = ", backend_config="
+ lidx = raw_line.find(backend_cfg_prefix)
+ if lidx >= 0:
+ backend_cfg_str = raw_line[lidx + len(backend_cfg_prefix) :]
+ # TODO: deal with exception raised.
+ backend_cfg = json.loads(backend_cfg_str)
+ raw_line = raw_line[:lidx]
+
+ # Clean the raw line.
+ raw_line = raw_line.rstrip()
+ size = len(raw_line)
+ raw_line = raw_line.lstrip()
+ indent = (size - len(raw_line)) // 2
+ return HloOperationInfo(raw_line, indent, metadata, backend_cfg)
+
+
+def parse_hlo_module(module: Lowered | Compiled) -> List[HloOperationInfo]:
+ """Parse an Hlo module, to be human-readable!
+
+ Note: `m.hlo_modules()[0].computations()[0].render_html()`
+ is also generating a nice HTML output!
+
+ Args:
+ module: HLO module or JAX stages compiled instance.
+ Returns:
+ List of HLO operation info.
+ """
+ assert isinstance(module, (Lowered, Compiled))
+ if isinstance(module, Lowered):
+ module = module.compile()
+ module_raw_txt = module.as_text()
+ module_lines = module_raw_txt.split("\n")
+ ops = [parse_hlo_operation_raw_line(line) for line in module_lines]
+ return ops
+
+
+def print_hlo_module(
+ module: Lowered | Compiled, metadata: bool = False, backend_cfg: bool = False, indent: int = 2
+) -> None:
+ """Human-readable Hlo module printing.
+
+ Args:
+ module: AOT Lowered or Compiled JAX module.
+ metadata: Print op metadata as well.
+ backend_cfg: Print op backend config as well.
+ """
+ cmds = parse_hlo_module(module)
+ for c in cmds:
+ print(c.as_text(metadata=metadata, backend_cfg=backend_cfg, indent=indent))

tests/utils/test_hlo.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2024 Graphcore Ltd. All rights reserved.
+import chex
+import jax
+import jax.numpy as jnp
+
+import jax_scalify as jsa
+
+
+class ScalifyHloUtilsTests(chex.TestCase):
+ def test__hlo_util__parse_hlo_module(self):
+ def matmul_fn(a, b):
+ return jax.lax.dot(a, b.T)
+
+ a = jax.core.ShapedArray((16, 48), jnp.float16)
+ b = jax.core.ShapedArray((32, 48), jnp.float16)
+
+ matmul_lowered = jax.jit(matmul_fn).lower(a, b)
+ matmul_compiled = matmul_lowered.compile()
+
+ ops = jsa.utils.parse_hlo_module(matmul_compiled)
+ assert len(ops) >= 6
+ # TODO: other tests???
+ # jsa.utils.print_hlo_module(matmul_compiled, backend_cfg=True, indent=2)