From 5fbdef572082948ca9c4aefffb9d09dbd50aedb8 Mon Sep 17 00:00:00 2001
From: illustris <rharikrishnan95@gmail.com>
Date: Sat, 18 Mar 2023 11:08:55 +0530
Subject: [PATCH] python3Packages.lightgbm: add GPU support

---
 .../python-modules/lightgbm/default.nix       | 23 +++++++++++++++----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/pkgs/development/python-modules/lightgbm/default.nix b/pkgs/development/python-modules/lightgbm/default.nix
index 1be852daa15..71604af7297 100644
--- a/pkgs/development/python-modules/lightgbm/default.nix
+++ b/pkgs/development/python-modules/lightgbm/default.nix
@@ -8,6 +8,11 @@
 , scikit-learn
 , llvmPackages ? null
 , pythonOlder
+, python
+, ocl-icd
+, opencl-headers
+, boost
+, gpuSupport ? true
 }:
 
 buildPythonPackage rec {
@@ -28,9 +33,13 @@ buildPythonPackage rec {
 
   dontUseCmakeConfigure = true;
 
-  buildInputs = lib.optionals stdenv.cc.isClang [
+  buildInputs = (lib.optionals stdenv.cc.isClang [
     llvmPackages.openmp
-  ];
+  ]) ++ (lib.optionals gpuSupport [
+    boost
+    ocl-icd
+    opencl-headers
+  ]);
 
   propagatedBuildInputs = [
     numpy
@@ -38,6 +47,10 @@ buildPythonPackage rec {
     scikit-learn
   ];
 
+  buildPhase = ''
+    ${python.pythonForBuild.interpreter} setup.py bdist_wheel ${lib.optionalString gpuSupport "--gpu"}
+  '';
+
   postConfigure = ''
     export HOME=$(mktemp -d)
   '';
@@ -51,11 +64,11 @@ buildPythonPackage rec {
     "lightgbm"
   ];
 
-  meta = with lib; {
+  meta = {
     description = "A fast, distributed, high performance gradient boosting (GBDT, GBRT, GBM or MART) framework";
     homepage = "https://github.com/Microsoft/LightGBM";
     changelog = "https://github.com/microsoft/LightGBM/releases/tag/v${version}";
-    license = licenses.mit;
-    maintainers = with maintainers; [ teh costrouc ];
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ teh costrouc ];
   };
 }