python3Packages.tvm-ffi: init at 0.19 (#318)

Author: danieldk

builder/lib/cache.nix

@@ -22,6 +22,7 @@
           ++ allOutputs build2cmake
           ++ allOutputs kernel-abi-check
           ++ allOutputs python3Packages.kernels
+          ++ allOutputs python3Packages.tvm-ffi
           ++ lib.optionals stdenv.hostPlatform.isLinux (allOutputs stdenvGlibc_2_27)
         );
       buildSetLinkFarm = buildSet: pkgs.linkFarm buildSet.torch.variant (buildSetOutputs buildSet);

flake.nix

@@ -169,6 +169,7 @@
               tabulate
               tomlkit
               torch
+              tvm-ffi
               types-pyyaml
               types-requests
               types-tabulate

nix/overlay.nix

@@ -109,6 +109,9 @@ in
           torchVersion = "2.9";
           xpuPackages = final.xpuPackages_2025_2;
         };
+
+        tvm-ffi = callPackage ./pkgs/python-modules/tvm-ffi {
+        };
       }
     )
     (import ./pkgs/python-modules/hooks)

nix/pkgs/python-modules/tvm-ffi/default.nix

@@ -0,0 +1,44 @@
+{
+  callPackage,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  cmake,
+  cython,
+  ninja,
+  python,
+  scikit-build-core,
+  setuptools-scm,
+}:
+
+buildPythonPackage rec {
+  pname = "tvm-ffi";
+  version = "0.1.9";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "apache";
+    repo = "tvm-ffi";
+    rev = "v${version}";
+    hash = "sha256-XnlM//WW2TbjbmzYBq6itJQ7R3J646UMVQUVhV5Afwc=";
+    fetchSubmodules = true;
+  };
+
+  build-system = [
+    cmake
+    cython
+    ninja
+    scikit-build-core
+    setuptools-scm
+  ];
+
+  dontUseCmakeConfigure = true;
+
+  postInstall = ''
+    ln -s $out/${python.sitePackages}/tvm_ffi/share $out/share
+  '';
+
+  passthru = callPackage ./variant.nix {
+    tvmFfiVersion = version;
+  };
+}

nix/pkgs/python-modules/tvm-ffi/variant.nix

@@ -0,0 +1,46 @@
+{
+  config,
+  cudaSupport ? config.cudaSupport,
+  metalSupport ? config.metalSupport or false,
+  rocmSupport ? config.rocmSupport,
+  xpuSupport ? config.xpuSupport or false,
+
+  cudaPackages,
+  rocmPackages,
+  xpuPackages,
+
+  lib,
+  stdenv,
+
+  tvmFfiVersion,
+}:
+
+let
+  flattenVersion =
+    version: lib.replaceStrings [ "." ] [ "" ] (lib.versions.majorMinor (lib.versions.pad 2 version));
+  backend =
+    if cudaSupport then
+      "cuda"
+    else if metalSupport then
+      "metal"
+    else if rocmSupport then
+      "rocm"
+    else if xpuSupport then
+      "xpu"
+    else
+      "cpu";
+  computeString =
+    if cudaSupport then
+      "cu${flattenVersion cudaPackages.cudaMajorMinorVersion}"
+    else if metalSupport then
+      "metal"
+    else if rocmSupport then
+      "rocm${flattenVersion (lib.versions.majorMinor rocmPackages.rocm.version)}"
+    else if xpuSupport then
+      "xpu${flattenVersion (lib.versions.majorMinor xpuPackages.oneapi-torch-dev.version)}"
+    else
+      "cpu";
+in
+{
+  variant = "tvm-ffi${flattenVersion (lib.versions.majorMinor tvmFfiVersion)}-${computeString}-${stdenv.hostPlatform.system}";
+}