nix-ascend/flake.nix

{
  description = "Ascend CANN: ascend-toolkit (CANN env), bisheng-wrapper (Nix bisheng; needs ASCEND_TOOLKIT_HOME), ascend-cann (both)";

  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-24.11";
    flake-utils.url = "github:numtide/flake-utils";
  };

  outputs = { self, nixpkgs, flake-utils }:
    flake-utils.lib.eachDefaultSystem (system:
      let
        pkgs = import nixpkgs { inherit system; };
        lib = pkgs.lib;

        bishengNixCflags =
          let
            gccCc = pkgs.gcc13.cc;
            cxxDev = lib.getDev gccCc;
            ccVersion = gccCc.version;
            host = pkgs.stdenv.hostPlatform.config;
          in
          builtins.concatStringsSep " " [
            "-isystem ${cxxDev}/include/c++/${ccVersion}"
            "-isystem ${cxxDev}/include/c++/${ccVersion}/${host}"
            "-isystem ${gccCc}/lib/gcc/${host}/${ccVersion}/include"
            "-isystem ${pkgs.glibc.dev}/include"
          ];

        bishengNixCflagsFile = pkgs.writeText "bisheng-nix-cflags" bishengNixCflags;

        fetchAscendToolkit = pkgs.fetchurl {
          url = "https://git.amadey.xyz/api/packages/mozart/generic/ascend-toolkit/8.5.0/ascend-toolkit-8.5.0-310p-full.tar.gz";
          sha256 = "sha256-VqC+bG0k4eZDuQkTkPKgCvigUQM8QJU6ByVvj77sk+Q=";
        };

        ascend-toolkit = pkgs.stdenvNoCC.mkDerivation rec {
          pname = "ascend-toolkit";
          version = "8.5.0";

          src = fetchAscendToolkit;

          nativeBuildInputs = [
            pkgs.patchelf
            pkgs.autoPatchelfHook
            pkgs.gcc.cc.lib
            pkgs.glibc
            pkgs.zlib
            pkgs.sqlite
          ];

          phases = [ "unpackPhase" "installPhase" "fixupPhase" ];

          unpackPhase = ''
            mkdir -p src-copy
            tar xzf "$src" -C src-copy
          '';

          installPhase = ''
            mkdir -p $out
            cp -r src-copy/* $out/
            chmod -R u+w $out/
            find $out -type f \( -name "bisheng" -o -name "bisheng-tune" \
              -o -name "atc" -o -name "atc.bin" -o -name "akt" \
              -o -name "ascendc_pack_kernel" -o -name "asc_dumper" \
              -o -name "asc_opc" -o -name "msprof" -o -name "ccec" \
              -o -name "cce-ld" -o -name "lld" -o -name "llvm-*" \
              -o -name "*.so" -o -name "*.so.*" \) -exec chmod +x {} +
            find $out/bin -type f -exec chmod +x {} + 2>/dev/null || true
            find $out/tools -type f \( -name "*.py" -o -name "*.sh" \) -exec chmod +x {} + 2>/dev/null || true
          '';

          setupHook = ./nix/ascend-toolkit-setup-hook.sh;

          postFixup = ''
            echo "Adding RPATH for toolkit-internal libraries..."
            local rpath="$out/lib64:$out/devlib:$out/tools/aml/lib64"

            for dir in bin tools/compiler x86_64-linux/bin x86_64-linux/ccec_compiler/bin \
                       x86_64-linux/bisheng_compiler/bin; do
              [ -d "$out/$dir" ] || continue
              find "$out/$dir" -type f -executable 2>/dev/null | while IFS= read -r f; do
                patchelf --add-rpath "$rpath" --force-rpath "$f" 2>/dev/null || true
              done
            done

            find $out -type f \( -name "*.so" -o -name "*.so.*" \) 2>/dev/null | while IFS= read -r f; do
              patchelf --add-rpath "$rpath" --force-rpath "$f" 2>/dev/null || true
            done

            local bisheng_config="$out/compiler/tikcpp/ascendc_kernel_cmake/legacy_modules/bisheng_config.cmake"
            if [ -f "$bisheng_config" ]; then
              cp ${./nix/bisheng_config_patched.cmake} "$bisheng_config"
            fi

            if [ -f "$out/nix-support/setup-hook" ]; then
              substituteInPlace "$out/nix-support/setup-hook" --subst-var-by ASCEND_TOOLKIT "$out"
            fi
          '';
        };

        bisheng-wrapper = pkgs.stdenv.mkDerivation {
          pname = "bisheng-wrapper";
          version = "8.5.0";

          dontUnpack = true;

          nativeBuildInputs = [ pkgs.makeWrapper ];

          installPhase = ''
            mkdir -p $out/bin $out/nix-support

            substituteAll ${./nix/bisheng-wrapper.sh} $out/bin/bisheng
            chmod +x $out/bin/bisheng

            substituteAll ${./nix/bisheng-wrapper.sh} $out/bin/bisheng-tune
            chmod +x $out/bin/bisheng-tune

            cat > $out/nix-support/add-flags.sh << 'EOF'
            : "''${NIX_CFLAGS_COMPILE_bisheng:=''${NIX_CFLAGS_COMPILE:-}}"
            : "''${NIX_CFLAGS_COMPILE_BEFORE_bisheng:=''${NIX_CFLAGS_COMPILE_BEFORE:-}}"
            : "''${NIX_CFLAGS_LINK_bisheng:=''${NIX_CFLAGS_LINK:-}}"
            : "''${NIX_LDFLAGS_bisheng:=''${NIX_LDFLAGS:-}}"
            export NIX_CC_WRAPPER_FLAGS_SET_bisheng=1
            EOF

            touch $out/nix-support/add-local-cc-cflags-before.sh

            cp ${bishengNixCflagsFile} $out/nix-support/bisheng-nix-cflags
            cp ${./nix/bisheng-wrapper-setup-hook.sh} $out/nix-support/setup-hook
            substituteInPlace $out/nix-support/setup-hook \
              --subst-var-by WRAPPER_OUT "$out"
          '';

          meta = with pkgs.lib; {
            description = "Bisheng wrapper with Nix CFLAGS; set ASCEND_TOOLKIT_HOME to your CANN toolkit (use with ascend-toolkit or any upstream install)";
            license = licenses.mit;
          };
        };

        # Single dependency for dev shells: propagates toolkit + bisheng-wrapper (hooks run in list order).
        ascend-cann = pkgs.stdenvNoCC.mkDerivation {
          pname = "ascend-cann";
          version = ascend-toolkit.version;

          dontUnpack = true;

          installPhase = "mkdir -p $out";

          propagatedBuildInputs = [
            ascend-toolkit
            bisheng-wrapper
          ];

          meta = with pkgs.lib; {
            description = "Meta package: CANN toolkit + Bisheng Nix wrapper (add only this, or use ascend-toolkit / bisheng-wrapper separately)";
            license = licenses.mit;
          };
        };
      in
      {
        packages = {
          inherit ascend-toolkit bisheng-wrapper ascend-cann;
          default = ascend-toolkit;
        };
      }
    );
}