ilgeco
108 lines
4.2 KiB
Python
108 lines
4.2 KiB
Python
"""Check reference navigation and parse examples with the release compiler."""
|
|
|
|
from html.parser import HTMLParser
|
|
import subprocess
|
|
import unittest
|
|
|
|
from check_reference import REFERENCE, ROOT, ReferenceParser
|
|
|
|
|
|
OPT = ROOT / "build_release/Release/bin/onnx-mlir-opt"
|
|
|
|
|
|
class IntroductionParser(HTMLParser):
|
|
def __init__(self, document):
|
|
super().__init__()
|
|
self.examples = {}
|
|
self.current = None
|
|
self.contents = []
|
|
self.in_contents = False
|
|
self.feed(document)
|
|
|
|
def handle_starttag(self, tag, attrs):
|
|
attrs = dict(attrs)
|
|
if tag == "pre" and "data-example" in attrs:
|
|
self.current = attrs["data-example"]
|
|
self.examples[self.current] = ""
|
|
if tag == "nav":
|
|
self.in_contents = attrs.get("aria-label") == "Table of contents"
|
|
if self.in_contents and tag == "a":
|
|
self.contents.append(attrs["href"])
|
|
|
|
def handle_data(self, data):
|
|
if self.current:
|
|
self.examples[self.current] += data
|
|
|
|
def handle_endtag(self, tag):
|
|
if tag == "pre":
|
|
self.current = None
|
|
if tag == "nav":
|
|
self.in_contents = False
|
|
|
|
|
|
class IntroductionTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.document = REFERENCE.read_text()
|
|
cls.introduction = IntroductionParser(cls.document)
|
|
|
|
def test_contents_covers_topics_and_indexes(self):
|
|
self.assertLess(
|
|
self.document.index('aria-label="Table of contents"'),
|
|
self.document.index('<section id="reading">'),
|
|
)
|
|
parser = ReferenceParser()
|
|
parser.feed(self.document)
|
|
for target in self.introduction.contents:
|
|
self.assertIn(target[1:], parser.ids)
|
|
self.assertEqual(set(self.introduction.contents), {
|
|
"#reading", "#compact", "#planning", "#compute-regions",
|
|
"#graph-batches", "#input-representations", "#scheduling",
|
|
"#destinations", "#memory-units", "#communication", "#contracts",
|
|
"#spatial-index", "#pim-index",
|
|
})
|
|
|
|
@unittest.skipUnless(OPT.is_file(), "Release onnx-mlir-opt is unavailable")
|
|
def test_examples_parse_and_print(self):
|
|
arguments = {
|
|
"addition": "%a: tensor<1x4xf32>, %b: tensor<1x4xf32>",
|
|
"rows": "%matrix: tensor<2x4xf32>",
|
|
"plan": "%a: tensor<1x2x3x4xf32>, %b: tensor<1x2x3x4xf32>",
|
|
"weighted": "%x: tensor<1x4xf32>",
|
|
"graph-batch": "%x: tensor<1x8xf32>",
|
|
"reconstruction": "%fragments: tensor<2x1x4xf32>",
|
|
"core-batch": None,
|
|
"destination": "%x: tensor<1x4xf32>, %dst: tensor<1x4xf32>",
|
|
"bufferized": "%x: memref<1x4xf32>, %dst: memref<1x4xf32>",
|
|
"copy": "%dst: memref<1x4xf32>, %src: memref<1x4xf32>",
|
|
"gather": "%dst: memref<1x4xf32>, %src: memref<1x8xf32>",
|
|
"send": "%x: tensor<1x4xf32>, %target: index",
|
|
"receive": "%dst: tensor<1x4xf32>, %source: index",
|
|
"signal": "%target: index, %event: index",
|
|
"wait": "%event: index, %count: index",
|
|
"deferred-group": "%a: tensor<1x4xf32>, %b: tensor<1x4xf32>",
|
|
}
|
|
self.assertEqual(set(self.introduction.examples), set(arguments))
|
|
for name, args in arguments.items():
|
|
with self.subTest(example=name):
|
|
snippet = self.introduction.examples[name]
|
|
if name in {"gather", "receive"}:
|
|
snippet = "%zero = arith.constant 0 : index\n" + snippet
|
|
if args is not None:
|
|
snippet = f"func.func @example({args}) {{\n{snippet}\nreturn\n}}"
|
|
program = f"module {{\n{snippet}\n}}\n"
|
|
command = [str(OPT), "--maccel=PIM"]
|
|
if name == "deferred-group":
|
|
command.append("--mlir-print-op-generic")
|
|
for _ in range(2):
|
|
result = subprocess.run(
|
|
command, input=program,
|
|
text=True, capture_output=True, timeout=30,
|
|
)
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
program = result.stdout
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|