Merge pull request #1987 from j2kun:pass-examples

PiperOrigin-RevId: 784751486

Author: copybara-github

.github/workflows/docs.yml

@@ -6,19 +6,61 @@ on:
     branches:
     - main
   pull_request:
-
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
 jobs:
+  check-cache:
+    runs-on:
+      labels: ubuntu-22.04-8core
+    outputs:
+      runner: ${{ steps.runner.outputs.runner }}
+    steps:
+      - name: Check out repository code
+        uses: actions/checkout@8e5e7e5ab8b370d6c329ec480221332ada57f0ab # pin@v3
+
+      - name: Cache bazel build artifacts
+        id: cache
+        uses: actions/cache@1bd1e32a3bdc45362d1e726936510720a7c30a57 # [email protected]
+        with:
+          path: |
+            ~/.cache/bazel
+          key: ${{ runner.os }}-${{ env.ImageVersion }}-bazel-${{ hashFiles('.bazelversion', '.bazelrc', 'WORKSPACE') }}-${{ hashFiles('bazel/import_llvm.bzl') }}
+          restore-keys: |
+            ${{ runner.os }}-${{ env.ImageVersion }}-bazel-${{ hashFiles('.bazelversion', '.bazelrc', 'WORKSPACE') }}-
+          lookup-only: true
+      - name: Select runner
+        id: runner
+        env:
+          CACHE_HIT: ${{ steps.cache.outputs.cache-hit  == 'true' }}
+        run: |
+          set -euo pipefail
+
+          if [[ "${CACHE_HIT}" == "true" ]]; then
+            echo "runner=ubuntu-22.04-8core" >> "$GITHUB_OUTPUT"
+          else
+            echo "runner=ubuntu-22.04-32core" >> "$GITHUB_OUTPUT"
+          fi
+
   build-and-deploy:
+    needs: check-cache
+    runs-on:
+      labels: ${{ needs.check-cache.outputs.runner }}
     permissions:
       contents: write
-    runs-on: ubuntu-22.04
-    concurrency:
-      group: ${{ github.workflow }}-${{ github.ref }}
     steps:
     - name: Checkout
       uses: actions/checkout@8e5e7e5ab8b370d6c329ec480221332ada57f0ab # pin@v3
       with:
         fetch-depth: 0         # Fetch all history for .GitInfo and .Lastmod
+    - name: Cache bazel build artifacts
+      uses: actions/cache@1bd1e32a3bdc45362d1e726936510720a7c30a57 # [email protected]
+      with:
+        path: |
+          ~/.cache/bazel
+        key: ${{ runner.os }}-${{ env.ImageVersion }}-bazel-${{ hashFiles('.bazelversion', '.bazelrc', 'WORKSPACE') }}-${{ hashFiles('bazel/import_llvm.bzl') }}
+        restore-keys: |
+          ${{ runner.os }}-${{ env.ImageVersion }}-bazel-${{ hashFiles('.bazelversion', '.bazelrc', 'WORKSPACE') }}-
 
     # This requires building mlir-tblgen, but may not require a full llvm build
     # as a result. It results in the files being added to their respective
@@ -32,16 +74,22 @@ jobs:
 
     - name: "Copy markdown files to docs/"
       run: |
-        python .github/workflows/copy_tblgen_files.py
+        python -m pip install --upgrade pip
+        python -m pip install pyyaml==6.0.2 fire==0.7.0
+        # heir-opt is needed to generate the doctest examples
+        bazel build //tools:heir-opt
+        python -m scripts.docs.copy_tblgen_files
 
-    # Please update the local install instructions at docs/README.md if changing hugo version
+    # Please update the local install instructions at docs/README.md if
+    # changing hugo version
     - name: Setup Hugo
       uses: peaceiris/actions-hugo@16361eb4acea8698b220b76c0d4e84e1fd22c61d # [email protected]
       with:
         hugo-version: '0.113.0'
         extended: true
 
-    # Please update the local install instructions at docs/README.md if changing node version
+    # Please update the local install instructions at docs/README.md if
+    # changing node version
     - name: Setup Node
       uses: actions/setup-node@0a44ba7841725637a19e28fa30b79a866c81b0a6 # v4.0.4
       with:

lib/Transforms/PolynomialApproximation/PolynomialApproximation.td

@@ -58,28 +58,7 @@ def PolynomialApproximation : Pass<"polynomial-approximation"> {
     These ops are replaced with `polynomial.eval` ops with a static polynomial
     attribute.
 
-    Examples:
-
-    ```mlir
-    %0 = math.exp %x {
-          degree = 3 : i32,
-          domain_lower = -1.0 : f64,
-          domain_upper = 1.0 : f64
-      } : f32
-    ```
-
-    is converted to
-
-    ```mlir
-    #ring_f64_ = #polynomial.ring<coefficientType = f64>
-    !poly = !polynomial.polynomial<ring = #ring_f64_>
-    %0 = polynomial.eval
-           #polynomial<typed_float_polynomial <
-             0.99458116404270657
-           + 0.99565537253615788x
-           + 0.54297028147256321x**2
-           + 0.17954582110873779x**3> : !poly>, %arg0 : f32
-    ```
+    (* example filepath=tests/Transforms/polynomial_approximation/doctest.mlir *)
   }];
   let dependentDialects = [
     "mlir::heir::polynomial::PolynomialDialect"

.github/workflows/copy_tblgen_files.py renamed to scripts/docs/copy_tblgen_files.py

@@ -4,6 +4,8 @@
 import yaml
 from pathlib import Path
 
+from scripts.docs.lit_to_doc import lit_to_doc
+
 SRC_BASE = "bazel-bin/lib"
 DEST_BASE = "docs/content/en/docs"
 PASSES_FILE = f"{DEST_BASE}/passes.md"
@@ -47,9 +49,10 @@ def rebuild_content(frontmatter, sorted_sections):
   frontmatter_str = (
       "---\n" + yaml.dump(frontmatter) + "---\n\n" if frontmatter else ""
   )
-  content_str = "".join(
-      [f"{header}\n{body}\n" for header, body in sorted_sections]
-  )
+  content_str = "".join([
+      f"{header.strip()}\n\n{body.strip()}\n\n"
+      for header, body in sorted_sections
+  ])
   return frontmatter_str + content_str
 
 
@@ -66,6 +69,25 @@ def sort_markdown_file_by_header(path):
     f.write(sorted_content)
 
 
+def generate_doctest_examples():
+  print("Generating doctests")
+  pattern = re.compile(r"^ *(\(\* example filepath=)(.*?)(\s*\*\))")
+
+  def replacer(match):
+    captured_filepath = match.group(2)
+    print(f"Running doctest for {captured_filepath}")
+    return lit_to_doc(lit_test_file=captured_filepath)
+
+  new_lines = []
+  with open(PASSES_FILE, "r") as passes_file:
+    for line in passes_file:
+      modified_line = pattern.sub(replacer, line)
+      new_lines.append(modified_line)
+
+  with open(PASSES_FILE, "w") as passes_file:
+    passes_file.write("".join(new_lines))
+
+
 if __name__ == "__main__":
   # Create passes.md file with the front matter
   with open(PASSES_FILE, "w") as f:
@@ -87,6 +109,7 @@ def sort_markdown_file_by_header(path):
         dest_file.write("\n")
 
   sort_markdown_file_by_header(PASSES_FILE)
+  generate_doctest_examples()
 
   print("Processing Dialects")
   Path(f"{DEST_BASE}/Dialects/").mkdir(parents=True, exist_ok=True)

scripts/docs/lit_to_doc.py

@@ -0,0 +1,86 @@
+import subprocess
+import os
+
+import fire
+
+from scripts.lit_to_bazel import (
+    get_command_without_bazel_prefix,
+    normalize_lit_test_file_arg,
+)
+
+
+def is_command_executable(command):
+  """Checks if a given command is executable on the system's PATH."""
+  path_env = os.environ.get("PATH")
+  treat_as_literal = os.path.sep in command or not path_env
+  if os.path.altsep:
+    treat_as_literal = treat_as_literal or os.path.altsep in command
+
+  if treat_as_literal:
+    return os.path.isfile(command) and os.access(command, os.X_OK)
+
+  for directory in path_env.split(os.pathsep):
+    full_path = os.path.join(directory, command)
+    if os.path.isfile(full_path) and os.access(full_path, os.X_OK):
+      return True
+  return False
+
+
+def get_doctest_input(lit_test_file) -> str:
+  """Return the mlir part of the given file, omitting lit/filecheck stuff."""
+  mlir_lines = []
+  with open(lit_test_file, "r") as f:
+    for line in f:
+      if "// RUN" in line or "// CHECK" in line:
+        continue
+      mlir_lines.append(line)
+
+  return "".join(mlir_lines).strip()
+
+
+def lit_to_doc(
+    lit_test_file: str,
+    git_root: str = "",
+):
+  """A helper CLI that converts MLIR test files to bazel run commands.
+
+  Args:
+    lit_test_file: The lit test file that should be converted to a bazel run
+      command.
+  """
+  lit_test_file = normalize_lit_test_file_arg(lit_test_file)
+  mlir_input = get_doctest_input(lit_test_file)
+  command = get_command_without_bazel_prefix(lit_test_file, git_root)
+  command = command.replace("%s", lit_test_file)
+
+  if not is_command_executable("heir-opt"):
+    command = command.replace("heir-opt", "bazel-bin/tools/heir-opt")
+
+  result = subprocess.run(
+      command, shell=True, capture_output=True, text=True, check=True
+  )
+  mlir_output = result.stdout.strip()
+
+  if not is_command_executable("heir-opt"):
+    command = command.replace("bazel-bin/tools/heir-opt", "heir-opt")
+
+  return f"""#### Example
+
+Command: `{command}`
+
+Input:
+
+```mlir
+{mlir_input}
+```
+
+Output:
+
+```mlir
+{mlir_output}
+```
+"""
+
+
+if __name__ == "__main__":
+  fire.Fire(lit_to_doc)

scripts/lit_to_bazel.py

@@ -83,19 +83,7 @@ def normalize_lit_test_file_arg(lit_test_file: str) -> str:
   return lit_test_file
 
 
-def lit_to_bazel(
-    lit_test_file: str,
-    git_root: str = "",
-    run: bool = False,
-    debug_dir: str = "/tmp/mlir",
-):
-  """A helper CLI that converts MLIR test files to bazel run commands.
-
-  Args:
-    lit_test_file: The lit test file that should be converted to a bazel run
-      command.
-  """
-
+def get_command_without_bazel_prefix(lit_test_file, git_root) -> str:
   if not git_root:
     git_root = pathlib.Path(__file__).parent.parent
     if not os.path.isdir(git_root / ".git"):
@@ -128,29 +116,48 @@ def lit_to_bazel(
     deduped_commands.append(command)
 
   joined = " ".join(deduped_commands)
+  return joined
+
+
+def lit_to_bazel(
+    lit_test_file: str,
+    git_root: str = "",
+    run: bool = False,
+    debug_dir: str = "/tmp/mlir",
+):
+  """A helper CLI that converts MLIR test files to bazel run commands.
+
+  Args:
+    lit_test_file: The lit test file that should be converted to a bazel run
+      command.
+  """
+
+  command = get_command_without_bazel_prefix(lit_test_file, git_root)
   # I would consider using bazel-bin/tools/heir-opt, but the yosys
   # requirement requires additional env vars to be set for the yosys and ABC
   # paths, which is not yet worth doing for this script.
-  joined = joined.replace(
+  command = command.replace(
       "heir-opt",
       "bazel run --noallow_analysis_cache_discard //tools:heir-opt --",
   )
-  joined = joined.replace(
+  command = command.replace(
       "heir-translate", f"{git_root}/bazel-bin/tools/heir-translate"
   )
-  joined = joined.replace("%s", str(pathlib.Path(lit_test_file).absolute()))
+  command = command.replace("%s", str(pathlib.Path(lit_test_file).absolute()))
 
   if run:
     # delete the debug dir if it exists
     if os.path.isdir(debug_dir):
       shutil.rmtree(debug_dir)
-    joined += f" --mlir-print-ir-after-all --mlir-print-ir-tree-dir={debug_dir}"
-    print(joined)
+    command += (
+        f" --mlir-print-ir-after-all --mlir-print-ir-tree-dir={debug_dir}"
+    )
+    print(command)
 
     # run the command
-    subprocess.run(joined, shell=True)
+    subprocess.run(command, shell=True)
   else:
-    print(joined)
+    print(command)
 
 
 if __name__ == "__main__":

tests/Transforms/polynomial_approximation/doctest.mlir

@@ -0,0 +1,13 @@
+// RUN: heir-opt --polynomial-approximation %s | FileCheck %s
+
+// CHECK: @test_exp
+func.func @test_exp(%x: f32) -> f32 {
+  // CHECK: polynomial.eval
+  // CHECK-SAME: x**3
+  // CHECK-NOT: x**4
+  %0 = math.exp %x {
+      degree = 3 : i32,
+      domain_lower = -1.0 : f64,
+      domain_upper = 1.0 : f64} : f32
+  return %0 : f32
+}