diff options
author | Weijia Wang <9713184+wegank@users.noreply.github.com> | 2023-11-18 00:33:42 +0100 |
---|---|---|
committer | GitHub <noreply@github.com> | 2023-11-18 00:33:42 +0100 |
commit | 4774c5365c67d337c6d7d860103466494910a728 (patch) | |
tree | dca9d7ffa6c94d7a807b88756d6f0434f0c807c4 /pkgs/by-name | |
parent | 9415631dfe1311648309c5cd357a734ddb0681d7 (diff) | |
parent | e3d4beac9df5b05d7b5edfcc74f02b49dc862145 (diff) |
Merge pull request #267209 from Madouura/pr/triton-llvm
openai-triton-llvm: fix aarch64 and cross-compilation
Diffstat (limited to 'pkgs/by-name')
-rw-r--r-- | pkgs/by-name/op/openai-triton-llvm/package.nix | 127 |
1 files changed, 89 insertions, 38 deletions
diff --git a/pkgs/by-name/op/openai-triton-llvm/package.nix b/pkgs/by-name/op/openai-triton-llvm/package.nix index 2fb56d0a63522..5744cfa10c11b 100644 --- a/pkgs/by-name/op/openai-triton-llvm/package.nix +++ b/pkgs/by-name/op/openai-triton-llvm/package.nix @@ -1,27 +1,49 @@ -{ config -, lib +{ lib , stdenv , fetchFromGitHub +, pkgsBuildBuild , pkg-config , cmake , ninja , git -, doxygen -, sphinx , libxml2 , libxcrypt , libedit , libffi +, libpfm , mpfr , zlib , ncurses +, doxygen +, sphinx +, which +, sysctl , python3Packages , buildDocs ? true , buildMan ? true , buildTests ? true +, llvmTargetsToBuild ? [ "NATIVE" ] # "NATIVE" resolves into x86 or aarch64 depending on stdenv +, llvmProjectsToBuild ? [ "llvm" "mlir" ] }: -stdenv.mkDerivation (finalAttrs: { +let + llvmNativeTarget = + if stdenv.hostPlatform.isx86_64 then "X86" + else if stdenv.hostPlatform.isAarch64 then "AArch64" + else throw "Currently unsupported LLVM platform '${stdenv.hostPlatform.config}'"; + + inferNativeTarget = t: if t == "NATIVE" then llvmNativeTarget else t; + llvmTargetsToBuild' = [ "AMDGPU" "NVPTX" ] ++ builtins.map inferNativeTarget llvmTargetsToBuild; + + # This LLVM version can't seem to find pygments/pyyaml, + # but a later update will likely fix this (openai-triton-2.1.0) + python = + if buildTests + then python3Packages.python.withPackages (p: with p; [ psutil pygments pyyaml ]) + else python3Packages.python; + + isNative = stdenv.hostPlatform == stdenv.buildPlatform; +in stdenv.mkDerivation (finalAttrs: { pname = "openai-triton-llvm"; version = "14.0.6-f28c006a5895"; @@ -33,7 +55,8 @@ stdenv.mkDerivation (finalAttrs: { "man" ]; - # See https://github.com/openai/triton/blob/main/python/setup.py and https://github.com/ptillet/triton-llvm-releases/releases + # See https://github.com/openai/triton/blob/main/python/setup.py + # and https://github.com/ptillet/triton-llvm-releases/releases src = fetchFromGitHub { owner = "llvm"; repo = "llvm-project"; @@ -46,7 +69,7 @@ stdenv.mkDerivation (finalAttrs: { cmake ninja git - python3Packages.python + python ] ++ lib.optionals (buildDocs || buildMan) [ doxygen sphinx @@ -58,6 +81,7 @@ stdenv.mkDerivation (finalAttrs: { libxcrypt libedit libffi + libpfm mpfr ]; @@ -69,37 +93,55 @@ stdenv.mkDerivation (finalAttrs: { sourceRoot = "${finalAttrs.src.name}/llvm"; cmakeFlags = [ - "-DLLVM_TARGETS_TO_BUILD=${ - let - # Targets can be found in - # https://github.com/llvm/llvm-project/tree/f28c006a5895fc0e329fe15fead81e37457cb1d1/clang/lib/Basic/Targets - # NOTE: Unsure of how "host" would function, especially given that we might be cross-compiling. - llvmTargets = [ "AMDGPU" "NVPTX" ] - ++ lib.optionals stdenv.isAarch64 [ "AArch64" ] - ++ lib.optionals stdenv.isx86_64 [ "X86" ]; - in - lib.concatStringsSep ";" llvmTargets - }" - "-DLLVM_ENABLE_PROJECTS=llvm;mlir" - "-DLLVM_INSTALL_UTILS=ON" - ] ++ lib.optionals (buildDocs || buildMan) [ - "-DLLVM_INCLUDE_DOCS=ON" - "-DMLIR_INCLUDE_DOCS=ON" - "-DLLVM_BUILD_DOCS=ON" - # "-DLLVM_ENABLE_DOXYGEN=ON" Way too slow, only uses one core - "-DLLVM_ENABLE_SPHINX=ON" - "-DSPHINX_OUTPUT_HTML=ON" - "-DSPHINX_OUTPUT_MAN=ON" - "-DSPHINX_WARNINGS_AS_ERRORS=OFF" - ] ++ lib.optionals buildTests [ - "-DLLVM_INCLUDE_TESTS=ON" - "-DMLIR_INCLUDE_TESTS=ON" - "-DLLVM_BUILD_TESTS=ON" - ]; + (lib.cmakeFeature "LLVM_TARGETS_TO_BUILD" (lib.concatStringsSep ";" llvmTargetsToBuild')) + (lib.cmakeFeature "LLVM_ENABLE_PROJECTS" (lib.concatStringsSep ";" llvmProjectsToBuild)) + (lib.cmakeFeature "LLVM_HOST_TRIPLE" stdenv.hostPlatform.config) + (lib.cmakeFeature "LLVM_DEFAULT_TARGET_TRIPLE" stdenv.hostPlatform.config) + (lib.cmakeBool "LLVM_INSTALL_UTILS" true) + (lib.cmakeBool "LLVM_INCLUDE_DOCS" (buildDocs || buildMan)) + (lib.cmakeBool "MLIR_INCLUDE_DOCS" (buildDocs || buildMan)) + (lib.cmakeBool "LLVM_BUILD_DOCS" (buildDocs || buildMan)) + # Way too slow, only uses one core + # (lib.cmakeBool "LLVM_ENABLE_DOXYGEN" (buildDocs || buildMan)) + (lib.cmakeBool "LLVM_ENABLE_SPHINX" (buildDocs || buildMan)) + (lib.cmakeBool "SPHINX_OUTPUT_HTML" buildDocs) + (lib.cmakeBool "SPHINX_OUTPUT_MAN" buildMan) + (lib.cmakeBool "SPHINX_WARNINGS_AS_ERRORS" false) + (lib.cmakeBool "LLVM_INCLUDE_TESTS" buildTests) + (lib.cmakeBool "MLIR_INCLUDE_TESTS" buildTests) + (lib.cmakeBool "LLVM_BUILD_TESTS" buildTests) + # Cross compilation code taken/modified from LLVM 16 derivation + ] ++ lib.optionals (!isNative) (let + nativeToolchainFlags = let + nativeCC = pkgsBuildBuild.targetPackages.stdenv.cc; + nativeBintools = nativeCC.bintools.bintools; + in [ + (lib.cmakeFeature "CMAKE_C_COMPILER" "${nativeCC}/bin/${nativeCC.targetPrefix}cc") + (lib.cmakeFeature "CMAKE_CXX_COMPILER" "${nativeCC}/bin/${nativeCC.targetPrefix}c++") + (lib.cmakeFeature "CMAKE_AR" "${nativeBintools}/bin/${nativeBintools.targetPrefix}ar") + (lib.cmakeFeature "CMAKE_STRIP" "${nativeBintools}/bin/${nativeBintools.targetPrefix}strip") + (lib.cmakeFeature "CMAKE_RANLIB" "${nativeBintools}/bin/${nativeBintools.targetPrefix}ranlib") + ]; + + # We need to repass the custom GNUInstallDirs values, otherwise CMake + # will choose them for us, leading to wrong results in llvm-config-native + nativeInstallFlags = [ + (lib.cmakeFeature "CMAKE_INSTALL_PREFIX" (placeholder "out")) + (lib.cmakeFeature "CMAKE_INSTALL_BINDIR" "${placeholder "out"}/bin") + (lib.cmakeFeature "CMAKE_INSTALL_INCLUDEDIR" "${placeholder "out"}/include") + (lib.cmakeFeature "CMAKE_INSTALL_LIBDIR" "${placeholder "out"}/lib") + (lib.cmakeFeature "CMAKE_INSTALL_LIBEXECDIR" "${placeholder "out"}/libexec") + ]; + in [ + (lib.cmakeBool "CMAKE_CROSSCOMPILING" true) + (lib.cmakeFeature "CROSS_TOOLCHAIN_FLAGS_NATIVE" (lib.concatStringsSep ";" + (lib.concatLists [ nativeToolchainFlags nativeInstallFlags ]))) + ]); postPatch = '' # `CMake Error: cannot write to file "/build/source/llvm/build/lib/cmake/mlir/MLIRTargets.cmake": Permission denied` chmod +w -R ../mlir + patchShebangs ../mlir/test/mlir-reduce # FileSystem permissions tests fail with various special bits rm test/tools/llvm-objcopy/ELF/mirror-permissions-unix.test @@ -107,9 +149,21 @@ stdenv.mkDerivation (finalAttrs: { substituteInPlace unittests/Support/CMakeLists.txt \ --replace "Path.cpp" "" + '' + lib.optionalString stdenv.isAarch64 '' + # Not sure why this fails + rm test/tools/llvm-exegesis/AArch64/latency-by-opcode-name.s + ''; + + postInstall = lib.optionalString (!isNative) '' + cp -a NATIVE/bin/llvm-config $out/bin/llvm-config-native ''; doCheck = buildTests; + + nativeCheckInputs = [ which ] + ++ lib.optionals stdenv.isDarwin [ sysctl ]; + + checkTarget = "check-all"; requiredSystemFeatures = [ "big-parallel" ]; meta = with lib; { @@ -117,9 +171,6 @@ stdenv.mkDerivation (finalAttrs: { homepage = "https://github.com/llvm/llvm-project"; license = with licenses; [ ncsa ]; maintainers = with maintainers; [ SomeoneSerge Madouura ]; - platforms = platforms.linux; - # Consider the derivation broken if we're not building for CUDA or ROCm, or if we're building for aarch64 - # and ROCm is enabled. See https://github.com/RadeonOpenCompute/ROCm/issues/1831#issuecomment-1278205344. - broken = stdenv.isAarch64 && !config.cudaSupport; + platforms = with platforms; aarch64 ++ x86; }; }) |