flake.nix

{
  description = "Apache TVM, machine learning optimization framework and compiler.";
  inputs.nixpkgs.url = "github:NixOS/nixpkgs/nixpkgs-unstable";
  inputs.flake-utils.url = "github:numtide/flake-utils";

  outputs = { self, nixpkgs, flake-utils }:
    flake-utils.lib.eachDefaultSystem (system:
      let
        pkgs = nixpkgs.legacyPackages.${system};
        nativeBuildInputs = [
          pkgs.bashInteractive
          pkgs.cmake
          pkgs.gtest
          pkgs.cargo
          pkgs.pkg-config
          pkgs.git
        ];
        buildInputs = [
          pkgs.llvm
          pkgs.python38
          pkgs.python38Packages.numpy
          pkgs.python38Packages.ipython
          pkgs.python38Packages.scipy
          pkgs.openssl
          pkgs.libbacktrace
        ];
        tvmFlag = name: enabled: "-DUSE_${name}=${if enabled then "ON" else "OFF"}";
      in
      rec {
        packages.tvm = with import nixpkgs { system = "${system}"; };
          stdenv.mkDerivation {
            inherit buildInputs;
            inherit nativeBuildInputs;
            name = "tvm";
            srcs = self;
            cmakeFlags = [
              (tvmFlag "LIBBACKTRACE" true)
              (tvmFlag "VTA" false)
              "-DINSTALL_DEV=ON"
            ];
            dontFixCmake = true;
          };
        defaultPackage = self.packages.${system}.tvm;
        #defaultPackage = packages.tvm;
        devShell = pkgs.mkShell {
          inherit nativeBuildInputs;
          inherit buildInputs;
        };
      });
}