nix-builder: switch to manylinux_2_28, dynamically link libstdc++ (#558)

Before this change, we were ensuring manylinux_2_28 compliance by:

* Building and linking against glibc 2.27 (there was no nixpkgs revision
  with 2.28).
* Linking libstdc++ statically.

Linking libstdc++ statically has turned out to be a pain point. If a
kernel uses libstdc++ functionality that does some global
initialization, the initialization of the dynamically-linked libstdc++
(e.g. through Torch) and the statically linked libstdc++ could conflict.
One example is `std::regex`, which indirectly initializes some global code
for locale handling.

As an alternative, I explored just linking libstdc++ dynamically, which
worked up to some point by avoiding certain C++ features. However gcc 13
moved the initialization of the standard stream objects into the shared
library. As a result, library compiled with gcc 13 or later requires a
libstdc++ that is newer than manylinux_2_28.

The EL8 (RHEL, AlmaLinux, etc.) gcc toolset used in the manylinux_2_28
solves this by using EL8 libstdc++ and providing a static library that
gets linked into a binary for newer C++ features. Since this is part of
a fairly large patchset to gcc, it does not seem feasible to reproduce
this easily in the Nix gcc derivations. So, instead, we repackage the
gcc toolsets from AlmaLinux as Nix derivations in this change. This
gives us a toolchain that is as close to the official manylinux_2_28
toolchain as possible.

The packaged toolchains are exposed as a standard nixpkgs stdenv, so
that they can be used with other Nix derivations (such as the
Torch/tvm-ffi extensions).

One exception is made in reproducing the toolchain and that is that we
build glibc 2.28 ourselves. glibc carries the dynamic loader. The
dynamic loader in AlmaLinux embeds FHS paths such as `/lib64`, which are
not valid in Nix, and lead to linking errors since the library directory
of e.g. `libc.so.6` cannot be found. By rebuilding glibc 2.28, we get a
dynamic loader with the correct paths into the Nix store.

Author: danieldk

.github/workflows/build_kernel.yaml

@@ -50,6 +50,7 @@ jobs:
           name: built-kernels-${{ matrix.arch }}
           path: |
             activation-kernel
+            cpp20-symbols-kernel
             cutlass-gemm-kernel
             cutlass-gemm-tvm-ffi-kernel
             extra-data

docs/source/_toctree.yml

@@ -66,3 +66,7 @@
     - local: builder-cli
       title: Builder CLI Reference
   title: CLI Reference
+- sections:
+    - local: builder/design-nix-builder
+      title: Nix Builder
+  title: Design

docs/source/builder/design-nix-builder.md

@@ -0,0 +1,115 @@
+# Nix Builder design
+
+## Introduction
+
+kernel-builder uses a Nix-based builder that orchestrates the build. The Nix
+builder provides:
+
+- Reproducible evaluation. The same Nix builder version will always produce
+  the same derivations (build recipes).
+- Largely reproducible builds by using a build sandbox that only has the
+  dependencies specified in a derivation.
+- Seamless creation of different build environments (e.g. different Torch
+  and CUDA combinations).
+
+## Kernel build steps
+
+A kernel derivation builds a kernel in the following steps:
+
+1. Generate CMake files for the kernel using
+   `kernel-builder create-pyproject`.
+2. Generate Ninja build files using CMake.
+3. Build the kernel using Ninja.
+4. Perform various checks on the compiled kernel, such as:
+   - Verify that the kernel only uses ABI3/`manylinux_2_28` symbols.
+   - Verify that the kernel can be loaded by the `kernels` Python package.
+5. Strip runpaths (ELF-embedded library directories) from kernel binaries
+   to make the kernel distribution-independent.
+
+## manylinux_2_28 compatibility
+
+To achieve `manylinux_2_28` compatibility, kernels are built using a
+toolchain similar to the `manylinux_2_28` Docker images. This toolchain
+is based on the gcc toolsets from AlmaLinux 8. `manylinux_2_28` [uses
+AlmaLinux 8 as its base](https://github.com/pypa/manylinux#manylinux_2_28-almalinux-8-based),
+so we have to compile against the same glibc/libstdc++ versions to
+ensure compatibility.
+
+We repackage the AlmaLinux 8 toolsets and libstdc++ as Nix derivations (see
+the `nix-builder/packages/manylinux_2_28` source directory). Then we merge
+various toolset packages to an unwrapped gcc that resembles unwrapped gcc in
+nixpkgs. Finally, we wrap binutils and gcc to combine them into a stdenv.
+
+The stdenv does not reuse glibc from AlmaLinux, since its dynamic loader has
+hardcoded FHS paths (`/lib64` etc.) that are not valid in Nix. Using this
+dynamic loader results in linking errors, since the paths in the dynamic
+loader are used as a last resort (to link glibc libraries). So, instead we
+build our own glibc 2.28 package
+(see `nix-builder/pkgs/manylinux_2_28/stdenv.nix`) and use that.
+
+## The package set pattern
+
+We repackage various existing package sets as Nix derivations. For instance,
+this is done for ROCm, XPU, and manylinux_2_28 packages. We do this because
+we want these libraries to be as close as what the user would install. This
+avoids compatibility issues between the kernels and the official vendor
+packages. For instance, suppose that we built a ROCm library as a shared
+library and ROCm provides the same library as a static library, then compiled
+kernels could use symbols that cannot be resolved when installing the official
+ROCm packages. Similarly, using the official packages allows us to test
+against the official upstram packages.
+
+These package sets all follow the same pattern:
+
+```nix
+{
+  lib,
+  callPackage,
+  newScope,
+  pkgs,
+}:
+
+{
+  packageMetadata,
+}:
+
+let
+  inherit (lib.fixedPoints) extends composeManyExtensions;
+
+  fixedPoint = final: {
+    inherit lib;
+  };
+  composed = lib.composeManyExtensions [
+    # Base package set.
+    (import ./components.nix { inherit packageMetadata; })
+
+    # Package-specific overrides.
+    (import ./overrides.nix)
+
+    # Additional overlays that extend the package set.
+    (import ./some-overlay.nix)
+  ];
+in
+lib.makeScope newScope (lib.extends composed fixedPoint)
+```
+
+We use a fixed point to build up the package set as a list of
+[overlays](https://nixos.org/manual/nixpkgs/stable/#sec-overlays-definition).
+This has various benefits. For instance, it allows us to refine the
+package set incrementally and we can refer to the final versions of
+packages in intermediate overlays.
+
+The package sets all use a similar list of overlays:
+
+- An initial overlay (`components.nix`) that applies a generic builder
+  to the package set metadata. The metadata typically comes from a Yum/DNF
+  repository that contains RPM packages.The generic builder will extract the
+  RPMs and move binaries, libraries, and headers to the right location. This
+  results in a set of Nix derivations that may or may not build.
+- The next overlay (`overrides.nix`) fixes up derivations generated by the
+  generic builder in the previous overlay that do not build. Fixing the
+  derivations typically consists of adding missing dependencies and changing
+  embedded FHS paths to Nix store paths.
+- Additional overlays with derivations that combine outputs from previous
+  overlays. One typical example are derivations that construct a full compiler
+  toolchain (e.g. `nix-builder/pkgs/manylinux_2_28/gcc-unwrapped.nix`).

examples/kernels/cpp20-symbols/build.toml

@@ -0,0 +1,16 @@
+[general]
+name = "cpp20-symbols"
+version = 1
+license = "Apache-2.0"
+backends = ["cpu"]
+
+[torch]
+src = [
+    "torch-ext/torch_binding.cpp",
+    "torch-ext/torch_binding.h",
+]
+
+[kernel.cpp20_symbols_cpu]
+backend = "cpu"
+depends = ["torch"]
+src = ["cpu/cpu.cpp"]

examples/kernels/cpp20-symbols/cpu/cpu.cpp

@@ -0,0 +1,18 @@
+#include <array>
+#include <charconv>
+#include <stdexcept>
+
+#include <torch/all.h>
+
+// std::to_chars(char*, char*, double) is a floating-point overload that
+// requires GLIBCXX_3.4.29, introduced in GCC 11. We use this to verify
+// that manylinux_2_28 kernels build correctly: the Red Hat toolset
+// statically links the newer libstdc++ symbols that exceed the system
+// GLIBCXX_3.4.25 ceiling of AlmaLinux 8 / RHEL 8.
+torch::Tensor float_to_chars(torch::Tensor const &input) {
+    std::array<char, 32> buf;
+    auto [ptr, ec] = std::to_chars(buf.begin(), buf.end(), input.item<double>());
+    if (ec != std::errc{})
+        throw std::runtime_error("to_chars failed");
+    return input;
+}

examples/kernels/cpp20-symbols/flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for invalid-cpp-manylinux-symbols kernel";
+
+  inputs = {
+    kernel-builder.url = "path:../../..";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

examples/kernels/cpp20-symbols/tests/__init__.py

@@ -0,0 +1,17 @@
+import cpp20_symbols
+import pytest
+import torch
+
+
+@pytest.mark.kernels_ci
+def test_float_to_chars_runs():
+    x = torch.tensor([3.14], dtype=torch.float64)
+    out = cpp20_symbols.float_to_chars(x)
+    torch.testing.assert_close(out, x)
+
+
+@pytest.mark.kernels_ci
+def test_float_to_chars_float32():
+    x = torch.tensor([2.71828], dtype=torch.float32)
+    out = cpp20_symbols.float_to_chars(x)
+    torch.testing.assert_close(out, x)

examples/kernels/cpp20-symbols/tests/test_cpp20_symbols.py

@@ -0,0 +1,10 @@
+import torch
+
+from ._ops import ops
+
+
+def float_to_chars(input: torch.Tensor) -> torch.Tensor:
+    return ops.float_to_chars(input)
+
+
+__all__ = ["float_to_chars"]

examples/kernels/cpp20-symbols/torch-ext/cpp20_symbols/__init__.py

@@ -0,0 +1,13 @@
+#include <torch/library.h>
+
+#include "registration.h"
+#include "torch_binding.h"
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+    ops.def("float_to_chars(Tensor input) -> Tensor");
+#if defined(CPU_KERNEL)
+    ops.impl("float_to_chars", torch::kCPU, &float_to_chars);
+#endif
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)