[llvm] [lit] Run builtin cat / diff in-process instead of spawning (PR #208024)

Prasoon Kumar via llvm-commits llvm-commits at lists.llvm.org
Thu Jul 30 20:41:28 PDT 2026


https://github.com/prasoon054 updated https://github.com/llvm/llvm-project/pull/208024

>From 08f36da60f1c8b04f228dcc1a4035aff77d00aa4 Mon Sep 17 00:00:00 2001
From: Prasoon Kumar <prasoonkumar054 at gmail.com>
Date: Tue, 7 Jul 2026 17:02:41 +0530
Subject: [PATCH] [lit] Run builtin cat/diff in-process instead of spawning

cat and diff are the only two builtins that still spawn a subprocess: every cat/diff on a RUN line spawns a fresh Python interpreter, which dominates wall time given how small lit's typical inputs are. Run them in-process instead.

The spawned-script path stays as a fallback for 'env VAR=... cat/diff' and 'not --crash cat/diff' for now. Removing it entirely is a follow-up PR. Output is byte-identical to the spawn path either way.

diff.py also switches its four output-encode sites from locale.getpreferredencoding to 'utf-8', fixing a UnicodeEncodeError that crashed Windows CI.

Signed-off-by: Prasoon Kumar <prasoonkumar054 at gmail.com>
---
 llvm/utils/lit/lit/ShellEnvironment.py      |  65 ++++
 llvm/utils/lit/lit/TestRunner.py            | 326 +++++++++++++++++++-
 llvm/utils/lit/lit/builtin_commands/diff.py |  18 +-
 3 files changed, 392 insertions(+), 17 deletions(-)

diff --git a/llvm/utils/lit/lit/ShellEnvironment.py b/llvm/utils/lit/lit/ShellEnvironment.py
index 1945865b19199..e2cf602de96ff 100644
--- a/llvm/utils/lit/lit/ShellEnvironment.py
+++ b/llvm/utils/lit/lit/ShellEnvironment.py
@@ -1,7 +1,11 @@
+from __future__ import annotations
+
+import io
 import os
 import platform
 import subprocess
 import tempfile
+from typing import BinaryIO, Protocol, TextIO
 
 import lit.util
 from lit.ShCommands import GlobItem
@@ -206,6 +210,67 @@ def processRedirects(cmd, stdin_source, cmd_shenv, opened_files):
     return std_fds
 
 
+def as_binary_reader(stream: None | int | io.TextIOBase | BinaryIO) -> BinaryIO:
+    """Adapts a builtin's stdin source into a binary, read()-able stream.
+
+    Args:
+        stream: Standard input source from pipeline dispatch. Supported types:
+            - None or subprocess sentinel (PIPE, DEVNULL, STDOUT) for no input
+            - Binary stream (BytesIO, temporary/spooled file, or 'rb' mode)
+            - Text stream (text-mode '<' redirect or upstream universal_newlines pipe)
+
+    Returns:
+        A binary, read()-able stream. Text streams are unwrapped to their
+        underlying buffer so an in-process builtin sees the same raw bytes a
+        child process would.
+    """
+    if stream is None or isinstance(stream, int):
+        # No real input to read.
+        return io.BytesIO(b"")
+    if isinstance(stream, io.TextIOBase):
+        buffer = getattr(stream, "buffer", None)
+        if buffer is not None:
+            return buffer
+        data = stream.read()
+        return io.BytesIO(data.encode() if isinstance(data, str) else data)
+    # Already a binary reader.
+    assert hasattr(stream, "read"), f"expected a binary reader, got {type(stream)!r}"
+    return stream
+
+
+class BinaryFileWriter:
+    """Writes bytes straight to a file descriptor, matching Popen's behavior.
+
+    Writes directly to the underlying file descriptor with os.write, bypassing
+    the wrapped file object's text-mode buffering and newline translation, so
+    the exact bytes passed in reach the fd unchanged.
+    """
+
+    # TODO: Replace __slots__ with @dataclass(slots=True)
+    # once the minimum Python version is bumped to 3.10
+    # https://github.com/llvm/llvm-project/issues/200531
+    __slots__ = ("fd",)
+
+    def __init__(self, fileobj: TextIO) -> None:
+        # Owned by opened_files, never closed here
+        self.fd = fileobj.fileno()
+
+    def write(self, data: bytes) -> int:
+        return os.write(self.fd, data)
+
+
+def binary_fd(fileobj: TextIO) -> BinaryFileWriter:
+    """Wraps a redirect file object as a byte-exact writer for in-process builtins."""
+    return BinaryFileWriter(fileobj)
+
+
+class ByteWriter(Protocol):
+    """Structural type for objects that support writing a bytes buffer."""
+
+    def write(self, data: bytes) -> object:
+        ...
+
+
 def expand_glob(arg, cwd):
     if isinstance(arg, GlobItem):
         return sorted(arg.resolve(cwd))
diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py
index 97cb54c63be64..681a56d9b45db 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import enum
+import io
 import os
 import pathlib
 import re
@@ -11,17 +12,23 @@
 import tempfile
 import threading
 import traceback
+from typing import IO, BinaryIO, Callable, List, TextIO, Union
 
 import lit.InprocBuiltins as InprocBuiltins
 import lit.ShUtil as ShUtil
 import lit.Test as Test
 import lit.util
+import lit.builtin_commands.cat as builtin_cat
+import lit.builtin_commands.diff as builtin_diff
 from lit.BooleanExpression import BooleanExpression
 from lit.ShCommands import Command
 from lit.ShellEnvironment import (
+    ByteWriter,
     InternalShellError,
     ShellCommandResult,
     ShellEnvironment,
+    as_binary_reader,
+    binary_fd,
     expand_glob,
     expand_glob_expressions,
     kAvoidDevNull,
@@ -212,6 +219,299 @@ def _replaceReadFile(match):
     return arguments
 
 
+class PipeIOConfig:
+    """Configuration for how one pipeline stage's output is routed, whether
+    to the next stage, a redirect file, or captured for the caller, and the
+    working directory the stage runs in.
+
+    Attributes:
+        out_sink: Where a non-last stage's output goes if the next stage
+            reads it. None for the last stage or a real redirect.
+        redirect_out: Binary writer for a '>'/'>>' target, else None.
+        capture_out: True if this stage's output should be captured and
+            returned to the caller instead of piped or redirected.
+        err_sink: Binary writer for a '2>file' target, else None.
+        cwd: The shell environment's current working directory.
+        merge_err: True for '2>&1', stderr writes to the out target instead.
+    """
+
+    # TODO: Replace __slots__ with @dataclass(slots=True)
+    # once the minimum Python version is bumped to 3.10
+    # https://github.com/llvm/llvm-project/issues/200531
+    __slots__ = (
+        "out_sink",
+        "redirect_out",
+        "capture_out",
+        "err_sink",
+        "cwd",
+        "merge_err",
+    )
+
+    def __init__(
+        self,
+        out_sink: IO[bytes] | None,
+        redirect_out: ByteWriter | None,
+        capture_out: bool,
+        err_sink: ByteWriter | None,
+        cwd: str,
+        merge_err: bool = False,
+    ) -> None:
+        self.out_sink = out_sink
+        self.redirect_out = redirect_out
+        self.capture_out = capture_out
+        self.err_sink = err_sink
+        self.cwd = cwd
+        self.merge_err = merge_err
+
+
+# A builtin's run(argv, stdin, stdout, stderr, cwd) entry point, e.g.
+# builtin_cat.run or builtin_diff.run. stdout/stderr accept anything with a
+# byte-oriented write() method: a real BytesIO or the duck-typed ByteWriter.
+RunFn = Callable[
+    [
+        List[str],
+        BinaryIO,
+        Union[IO[bytes], ByteWriter, None],
+        Union[IO[bytes], ByteWriter, None],
+        str,
+    ],
+    int,
+]
+
+
+class InProcessPipe:
+    """Popen-compatible shim for an I/O-heavy builtin command.
+
+    Runs the builtin (e.g., cat or diff) in-process instead of spawning
+    it. Anything that isn't a lit builtin still spawns a real subprocess.
+
+    Exposes the same interface as subprocess.Popen, so it's a drop-in
+    replacement wherever a real subprocess is expected.
+
+    Attributes:
+        returncode: The builtin's exit code.
+        stdout: The captured stdout if this stage's output was captured,
+            else None.
+        stderr: The captured stderr.
+    """
+
+    # TODO: Replace __slots__ with @dataclass(slots=True)
+    # once the minimum Python version is bumped to 3.10
+    # https://github.com/llvm/llvm-project/issues/200531
+    __slots__ = ("returncode", "stdout", "stderr", "_out", "_err")
+
+    def __init__(
+        self,
+        run_fn: RunFn,
+        args: List[str],
+        stdin: None | int | io.TextIOBase | BinaryIO,
+        stage_io: PipeIOConfig,
+    ) -> None:
+        """Runs run_fn synchronously and captures the result.
+
+        Args:
+            run_fn: The builtin's run function, e.g. cat.run or diff.run.
+            args: The argv to pass to run_fn, args[0] is the command name.
+            stdin: This stage's input stream.
+            stage_io: Where this stage's output and error streams go, and
+                its cwd.
+        """
+        in_stream = as_binary_reader(stdin)
+        out_buf = io.BytesIO() if stage_io.capture_out else None
+        out_target = (
+            out_buf
+            if stage_io.capture_out
+            else (stage_io.redirect_out or stage_io.out_sink)
+        )
+        if stage_io.merge_err:
+            err_target, err_buf = out_target, None
+        elif stage_io.err_sink is not None:
+            err_target, err_buf = stage_io.err_sink, None
+        else:
+            err_buf = io.BytesIO()
+            err_target = err_buf
+
+        self.returncode = run_fn(args, in_stream, out_target, err_target, stage_io.cwd)
+
+        self._out = out_buf.getvalue() if out_buf is not None else b""
+        self._err = err_buf.getvalue() if err_buf is not None else b""
+        self.stdout = io.BytesIO(self._out) if stage_io.capture_out else None
+        self.stderr = io.BytesIO(self._err)
+
+    def communicate(self) -> tuple[bytes, bytes]:
+        return (self._out, self._err)
+
+    def wait(self) -> int:
+        return self.returncode
+
+    def poll(self) -> int:
+        return self.returncode
+
+    def kill(self) -> None:
+        pass
+
+    def terminate(self) -> None:
+        pass
+
+
+def _resolve_redirect_out(stdout: int | TextIO) -> ByteWriter | None:
+    """Wraps a '>' or '>>' redirect target as a binary writer.
+
+    Args:
+        stdout: This stage's stdout target, a subprocess sentinel
+            (PIPE/STDOUT) or an open file object for a real redirect.
+
+    Returns:
+        A binary writer wrapping the redirect file. None if stdout is a
+        PIPE or STDOUT sentinel instead of a real redirect file, since
+        that case doesn't need this sink.
+    """
+    if isinstance(stdout, int):
+        return None
+    return binary_fd(stdout)
+
+
+def _should_capture(is_last: bool, stdout: int | TextIO) -> bool:
+    """Whether this stage's output should be captured for the caller.
+
+    Args:
+        is_last: True for the pipeline's final stage, the only stage allowed
+            to capture instead of feeding the next stage or a redirect.
+        stdout: The stage's stdout target. Capture only applies when
+            this is subprocess.PIPE, ruling out a real redirect file.
+
+    Returns:
+        True only for the last stage of a piped command reading from a
+        PIPE, the one stage that reports output back to the caller
+        instead of writing to a temp file or redirect.
+    """
+    return is_last and stdout == subprocess.PIPE
+
+
+def _make_out_sink(stdout: int | TextIO, is_last: bool) -> IO[bytes] | None:
+    """Builds the output sink for a non-last stage feeding a pipe.
+
+    Args:
+        stdout: The stage's stdout target. Only subprocess.PIPE gets a
+            sink here, a real redirect target writes there directly.
+        is_last: True for the pipeline's final stage, which streams to
+            the caller instead of spooling for a downstream stage.
+
+    Returns:
+        A SpooledTemporaryFile, which stays in memory unless its size or
+        a downstream call to .fileno() forces it to disk. None for the
+        last stage or a non-pipe target, which don't need this sink.
+    """
+    if stdout == subprocess.PIPE and not is_last:
+        return tempfile.SpooledTemporaryFile(max_size=1 << 20)
+    return None
+
+
+def _resolve_err_sink(stderr: int | TextIO, merge_err: bool) -> ByteWriter | None:
+    """Wraps a '2>file' redirect target as a binary writer.
+
+    Args:
+        stderr: This stage's stderr target. Only a real redirect file
+            needs wrapping here, PIPE/STDOUT are handled elsewhere.
+        merge_err: Whether stderr is merged into stdout (2>&1).
+
+    Returns:
+        A binary writer wrapping the redirect file. None when stderr is
+        merged into stdout or is a PIPE/STDOUT sentinel, since both cases
+        are captured into the result instead of written through a sink
+        here.
+    """
+    if merge_err or isinstance(stderr, int):
+        return None
+    return binary_fd(stderr)
+
+
+def _should_run_inproc(
+    builtin_fn: RunFn | None,
+    not_crash: bool,
+    cmd_shenv: ShellEnvironment,
+    shenv: ShellEnvironment,
+) -> bool:
+    """Whether this pipeline stage can run in-process instead of spawning.
+
+    Args:
+        builtin_fn: The builtin's run function for this command (e.g.
+            builtin_cat.run), or None if the command has no in-process
+            implementation.
+        not_crash: True if the command is wrapped in 'not --crash'.
+        cmd_shenv: The environment this specific command runs under.
+        shenv: The pipeline's shared environment.
+
+    Returns:
+        False if builtin_fn is None, if not_crash is set, or if cmd_shenv
+        is not shenv (a per-command 'env' only takes effect on a spawned
+        child, not an in-process call sharing our environment). True
+        otherwise.
+    """
+    return builtin_fn is not None and not not_crash and cmd_shenv is shenv
+
+
+def _run_inproc_stage(
+    args: List[str],
+    builtin_fn: RunFn,
+    stdin: None | int | io.TextIOBase | BinaryIO,
+    stdout: int | TextIO,
+    stderr: int | TextIO,
+    cmd_shenv: ShellEnvironment,
+    is_last: bool,
+    named_temp_files: List[str],
+) -> tuple[InProcessPipe, IO[bytes] | int]:
+    """Runs one in-process pipeline stage and returns its default_stdin.
+
+    Args:
+        args: The command's argument vector.
+        builtin_fn: The builtin's run function to execute (e.g.,
+            cat.run or diff.run).
+        stdin: Current stage's input. None or a placeholder int means no
+            real input, otherwise an open binary or text stream to read
+            from.
+        stdout: Target for the stage's stdout, either subprocess.PIPE or
+            an open redirect file. Whether PIPE feeds the next stage or
+            is captured depends on is_last.
+        stderr: Target for the stage's stderr. Set to subprocess.STDOUT
+            when merged into stdout.
+        cmd_shenv: Shell environment the stage runs under.
+        is_last: True for the final stage in the pipeline. Together with
+            stdout being subprocess.PIPE, this determines whether output
+            is captured instead of feeding a downstream stage.
+        named_temp_files: Caller-owned cleanup list. Any temp file created
+            to stand in for /dev/null is appended here.
+
+    Returns:
+        The InProcessPipe wrapping the stage, and the stdin the next stage
+        should read from (the spooled output, or subprocess.PIPE).
+    """
+    if kAvoidDevNull:
+        for arg_idx, arg in enumerate(args):
+            if isinstance(arg, str) and kDevNull in arg:
+                devnull = tempfile.NamedTemporaryFile(delete=False)
+                devnull.close()
+                named_temp_files.append(devnull.name)
+                args[arg_idx] = arg.replace(kDevNull, devnull.name)
+    args = expand_glob_expressions(args, cmd_shenv.cwd)
+    merge_err = stderr == subprocess.STDOUT
+    stage_io = PipeIOConfig(
+        out_sink=_make_out_sink(stdout, is_last),
+        redirect_out=_resolve_redirect_out(stdout),
+        capture_out=_should_capture(is_last, stdout),
+        err_sink=_resolve_err_sink(stderr, merge_err),
+        cwd=cmd_shenv.cwd,
+        merge_err=merge_err,
+    )
+    proc = InProcessPipe(builtin_fn, args, stdin, stage_io)
+    if stage_io.out_sink is not None:
+        stage_io.out_sink.seek(0)
+        default_stdin = stage_io.out_sink
+    else:
+        default_stdin = subprocess.PIPE
+    return proc, default_stdin
+
+
 def _executeShCmd(cmd, shenv, results, timeoutHelper):
     if timeoutHelper.timeoutReached():
         # Prevent further recursion if the timeout has been hit
@@ -268,6 +568,10 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper):
         "umask": InprocBuiltins.executeBuiltinUmask,
         ":": InprocBuiltins.executeBuiltinColon,
     }
+    pipeline_builtins = {
+        "cat": builtin_cat.run,
+        "diff": builtin_diff.run,
+    }
     # To avoid deadlock, we use a single stderr stream for piped
     # output. This is null until we have seen some output using
     # stderr.
@@ -362,13 +666,14 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper):
             results.append(result)
             return result.exitCode
 
-        # Resolve any out-of-process builtin command before adding back 'not'
-        # commands.
-        if args[0] in builtin_commands:
+        builtin_fn = pipeline_builtins.get(args[0])
+        use_inproc = _should_run_inproc(builtin_fn, not_crash, cmd_shenv, shenv)
+        if not use_inproc and args[0] in builtin_commands:
             args.insert(0, sys.executable)
             cmd_shenv.env["PYTHONPATH"] = os.path.dirname(os.path.abspath(__file__))
             args[1] = os.path.join(builtin_commands_dir, args[1] + ".py")
 
+
         # We had to search through the 'not' commands to find all the 'env'
         # commands and any other in-process builtin command.  We don't want to
         # reimplement 'not' and its '--crash' here, so just push all 'not'
@@ -395,6 +700,21 @@ def _executeShCmd(cmd, shenv, results, timeoutHelper):
             j, default_stdin, cmd_shenv, opened_files
         )
 
+        if use_inproc:
+            proc, default_stdin = _run_inproc_stage(
+                args,
+                builtin_fn,
+                stdin,
+                stdout,
+                stderr,
+                cmd_shenv,
+                is_last=j is cmd.commands[-1],
+                named_temp_files=named_temp_files,
+            )
+            procs.append(proc)
+            proc_not_counts.append(not_count)
+            proc_not_fail_if_crash.append(False)
+            continue
         # If stderr wants to come from stdout, but stdout isn't a pipe, then put
         # stderr on a pipe and treat it as stdout.
         if stderr == subprocess.STDOUT and stdout != subprocess.PIPE:
diff --git a/llvm/utils/lit/lit/builtin_commands/diff.py b/llvm/utils/lit/lit/builtin_commands/diff.py
index 9d1a398106664..3a85d2210705f 100644
--- a/llvm/utils/lit/lit/builtin_commands/diff.py
+++ b/llvm/utils/lit/lit/builtin_commands/diff.py
@@ -92,11 +92,7 @@ def compareTwoBinaryFiles(flags, filepaths, filelines, stdout):
     )
 
     for diff in diffs:
-        stdout.write(
-            diff.decode(errors="backslashreplace").encode(
-                locale.getpreferredencoding(False)
-            )
-        )
+        stdout.write(diff.decode(errors="backslashreplace").encode("utf-8"))
         exitCode = 1
     return exitCode
 
@@ -153,9 +149,7 @@ def printDirVsFile(dir_path, file_path, stdout):
         msg = "File %s is a directory while file %s is a regular file"
     else:
         msg = "File %s is a directory while file %s is a regular empty file"
-    stdout.write(
-        (msg % (dir_path, file_path) + "\n").encode(locale.getpreferredencoding(False))
-    )
+    stdout.write((msg % (dir_path, file_path) + "\n").encode("utf-8"))
 
 
 def printFileVsDir(file_path, dir_path, stdout):
@@ -163,16 +157,12 @@ def printFileVsDir(file_path, dir_path, stdout):
         msg = "File %s is a regular file while file %s is a directory"
     else:
         msg = "File %s is a regular empty file while file %s is a directory"
-    stdout.write(
-        (msg % (file_path, dir_path) + "\n").encode(locale.getpreferredencoding(False))
-    )
+    stdout.write((msg % (file_path, dir_path) + "\n").encode("utf-8"))
 
 
 def printOnlyIn(basedir, path, name, stdout):
     stdout.write(
-        ("Only in %s: %s\n" % (os.path.join(basedir, path), name)).encode(
-            locale.getpreferredencoding(False)
-        )
+        ("Only in %s: %s\n" % (os.path.join(basedir, path), name)).encode("utf-8")
     )
 
 



More information about the llvm-commits mailing list