[Mlir-commits] [mlir] [mlir][pygments] Fix op-name highlighting, add low-precision float types (PR #220532)
Perry Gibson
llvmlistbot at llvm.org
Wed Sep 2 02:54:30 PDT 2026
https://github.com/Wheest created https://github.com/llvm/llvm-project/pull/220532
Two fixes to the MLIR Pygments lexer, plus a small test.
### Operation names lose their highlighting when indented
The `\s+` rule in the `root` state, added in #206194 to stop Pygments warning about unmatched whitespace, sits before the `^(\s*)(name)` rules that give operation names `Name.Builtin`. It consumes the indent, so those anchored rules can no longer match. On a four-line function only the operation at line 1 column 0 keeps its highlighting:
```mlir
func.func @f() {
%0 = arith.constant 1 : i32
return
}
```
lexes `func.func` as `Name.Builtin` but `arith.constant` and `return` as `Name.Other`, and the `=` degrades from `Operator` to `Punctuation`.
This moves the rule after the two operation-name rules and splits it at the newline. The split matters: a single `\s+` that swallows a newline plus the next line's indent leaves the position mid-line, which reintroduces the same problem one line later.
#206194's motivation is preserved. Its own example, ` %0 = cir.foo : !cir.int<s, 32>`, still lexes with zero `Error` tokens, so the docs builds it fixed stay green — and `cir.foo` is now highlighted, which it was not before.
### The low-precision builtin float types lex as `Error`
The builtin-type rule covers `bf16` through `f128`, so the eleven low-precision floats in `BuiltinTypes.td` — the float8 family, the OCP MX element formats `f6E2M3FN`/`f6E3M2FN`/`f4E2M1FN`, and the `f8E8M0FNU` scale type — plus `tf32` fall through it. Inside `tensor<1x6xf8E8M0FNU>` the `x` dimension separator abuts the element type, so no `\b`-anchored rule can match and the name breaks up into a mix of `Number` and `Error` tokens.
Written as one generative alternative so a new element format needs no further change, and applied to both the `root` and `angled-type` rules, whose `NOTE` comments ask that they be kept in sync. Per the README's "known to be incomplete and support mostly core IR with a subset of built-in types. Additions and customizations are welcome."
### Test
`mlir/utils/pygments/test_mlir_lexer.py`, stdlib `unittest`, runnable as `python -m unittest` from that directory. Two cases: an indented operation keeps `Name.Builtin`, and each of the 18 builtin float types is one token both bare and as a tensor element type. It fails 25 times against the current lexer and passes with this change.
It is deliberately **not** wired into `check-mlir`: nothing under `mlir/utils` is, and Pygments is only pinned in `llvm/docs/requirements.txt`, so a lit test would put an optional Python dependency in front of a check-mlir target. Happy to add one with a `pygments` lit feature probe instead if that is preferred.
Draft while I get feedback on the test's placement and whether it is wanted at all.
>From 9496cfa44fb49c0fdd61a9f3896ccb364c4a655e Mon Sep 17 00:00:00 2001
From: Perry Gibson <perry at fractile.ai>
Date: Wed, 2 Sep 2026 10:38:39 +0100
Subject: [PATCH] [mlir][pygments] Fix op-name highlighting, add low-precision
float types
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
The root state's `\s+` rule, added in #206194 to stop Pygments emitting
warnings for unmatched whitespace, is matched before the `^(\s*)(name)`
rules that highlight operation names. It consumes the indent, so those
anchored rules can no longer match and every indented operation loses its
Name.Builtin token — on a four-line function only the operation at line 1
column 0 survives.
Move the rule after them and split it at the newline, since a single `\s+`
that swallows a newline plus the next line's indent leaves the position
mid-line and reintroduces the same problem. #206194's own example,
` %0 = cir.foo : !cir.int<s, 32>`, still lexes with no Error tokens, so
the docs builds it fixed stay green, and `cir.foo` is now highlighted.
Also add the eleven low-precision builtin float types from BuiltinTypes.td
(the float8 family, the OCP MX element formats f6E2M3FN/f6E3M2FN/f4E2M1FN,
and the f8E8M0FNU scale type) plus tf32, none of which the builtin-type
rule covered. Inside `tensor<1x6xf8E8M0FNU>` the `x` dimension separator
abuts the element type, so no \b-anchored rule matches and the name lexed
as a mix of Number and Error tokens. Written as one generative alternative
so a new element format needs no further change, and applied to both the
root and angled-type rules, which the NOTE comments ask to keep in sync.
Adds a unittest-based test for the lexer, runnable from the directory as
`python -m unittest`. It is not wired into check-mlir: nothing under
mlir/utils is, and Pygments is only pinned in llvm/docs/requirements.txt,
so making it a lit test would put an optional dependency in front of a
check-mlir target.
---
mlir/utils/pygments/mlir_lexer.py | 18 +++++++++--
mlir/utils/pygments/test_mlir_lexer.py | 45 ++++++++++++++++++++++++++
2 files changed, 60 insertions(+), 3 deletions(-)
create mode 100644 mlir/utils/pygments/test_mlir_lexer.py
diff --git a/mlir/utils/pygments/mlir_lexer.py b/mlir/utils/pygments/mlir_lexer.py
index c493d20272cc6..985ae32a7fc31 100644
--- a/mlir/utils/pygments/mlir_lexer.py
+++ b/mlir/utils/pygments/mlir_lexer.py
@@ -42,7 +42,6 @@ class VariableList(RegexLexer):
tokens = {
"root": [
- (r"\s+", Text),
# Comments
(r"//.*?$", Comment.Single),
# operation name with assignment: %... = op.name
@@ -52,6 +51,12 @@ class VariableList(RegexLexer):
),
# operation name without result
(r"^(\s*)([A-Za-z0-9_\.\$\-]+)\b(?=[^<:])", bygroups(Text, Name.Builtin)),
+ # Whitespace. Must come after the anchored operation-name rules
+ # above and be split at the newline: a single `\s+` swallowing a
+ # newline plus the next line's indent leaves the position
+ # mid-line, so `^`-anchored rules never fire again.
+ (r"\n", Text),
+ (r"[ \t]+", Text),
# Attribute alias definition: #name =
(
r"^(\s*)(#[_A-Za-z0-9\$\-\.]+)(\b)(\s*=)",
@@ -73,7 +78,12 @@ class VariableList(RegexLexer):
# types by exclamation or builtin names
(r"![_A-Za-z0-9\$\-\.]+\b", Keyword.Type),
# NOTE: please sync changes to corresponding builtin type rule in "angled-type"
- (r"\b(bf16|f16|f32|f64|f80|f128|index|none|(u|s)?i[0-9]+)\b", Keyword.Type),
+ (
+ r"\b(bf16|f16|f32|f64|f80|f128|tf32|index|none"
+ r"|f[0-9]+E[0-9]+M[0-9]+(B[0-9]+)?(FN(UZ|U)?)?"
+ r"|(u|s)?i[0-9]+)\b",
+ Keyword.Type,
+ ),
# container-like dialect types (tensor<...>, memref<...>, vector<...>)
(
r"\b(complex|memref|tensor|tuple|vector)\s*(<)",
@@ -121,7 +131,9 @@ class VariableList(RegexLexer):
# element / builtin types inside angle brackets (no word-boundary)
# NOTE: please sync changes to corresponding builtin type rule in "root"
(
- r"(?:bf16|f16|f32|f64|f80|f128|index|none|(?:[us]?i[0-9]+))",
+ r"(?:bf16|f16|f32|f64|f80|f128|tf32|index|none"
+ r"|f[0-9]+E[0-9]+M[0-9]+(?:B[0-9]+)?(?:FN(?:UZ|U)?)?"
+ r"|(?:[us]?i[0-9]+))",
Keyword.Type,
),
# also allow nested container-like types to be recognized
diff --git a/mlir/utils/pygments/test_mlir_lexer.py b/mlir/utils/pygments/test_mlir_lexer.py
new file mode 100644
index 0000000000000..fb0bd7ddd2cfb
--- /dev/null
+++ b/mlir/utils/pygments/test_mlir_lexer.py
@@ -0,0 +1,45 @@
+# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+"""Tests for mlir_lexer.py. Run with `python -m unittest` from this directory.
+
+Requires Pygments (`pip install Pygments`), as the lexer itself does.
+"""
+
+import unittest
+
+from mlir_lexer import MlirLexer
+from pygments.token import Error, Keyword, Name
+
+# Every builtin float type keyword, per mlir/include/mlir/IR/BuiltinTypes.td.
+FLOAT_TYPES = "bf16 f16 f32 f64 f80 f128 tf32 f8E3M4 f8E4M3 f8E4M3B11FNUZ \
+f8E4M3FN f8E4M3FNUZ f8E5M2 f8E5M2FNUZ f8E8M0FNU f6E2M3FN f6E3M2FN f4E2M1FN".split()
+
+
+def lex(source):
+ """Lex source, dropping whitespace-only tokens."""
+ return [(t, v) for t, v in MlirLexer().get_tokens(source) if v.strip()]
+
+
+class TestMlirLexer(unittest.TestCase):
+ def test_indented_operation_names(self):
+ """Whitespace rules must not consume the indent the op-name rules anchor on."""
+ got = lex("func.func @f() {\n %0 = arith.constant 1 : i32\n return\n}\n")
+ self.assertEqual(
+ [v for t, v in got if t is Name.Builtin],
+ ["func.func", "arith.constant", "return"],
+ )
+
+ def test_builtin_float_types(self):
+ """Each float type is one token, bare and as a tensor element type."""
+ for name in FLOAT_TYPES:
+ for source in ("%%0 = a.b %%c : %s\n", "%%0 = a.b %%c : tensor<4x8x%s>\n"):
+ with self.subTest(type=name, source=source):
+ got = lex(source % name)
+ self.assertIn((Keyword.Type, name), got)
+ self.assertEqual([v for t, v in got if t is Error], [])
+
+
+if __name__ == "__main__":
+ unittest.main()
More information about the Mlir-commits
mailing list