diff --git a/kmath-noa/README.md b/kmath-noa/README.md index 9ca1933f6..9a629338b 100644 --- a/kmath-noa/README.md +++ b/kmath-noa/README.md @@ -18,14 +18,8 @@ the [GNU](https://gcc.gnu.org/) toolchain. For `GPU` kernels, we require a compa installation. If you are on Windows, we recommend setting up everything on [WSL](https://docs.nvidia.com/cuda/wsl-user-guide/index.html). -To install the library, you have to publish -locally `kmath-core`, `kmath-tensors` with `kmath-noa`: - +To install the library, simply publish it locally: ``` -./gradlew -q :kmath-core:publishToMavenLocal :kmath-tensors:publishToMavenLocal :kmath-noa:publishToMavenLocal +./gradlew -q :kmath-noa:publishToMavenLocal ``` -This builds `jtorch` a JNI wrapper for `NOA/LibTorch`, placed inside: - -`~/.konan/third-party/kmath-noa-/cpp-build` - diff --git a/kmath-noa/build.gradle.kts b/kmath-noa/build.gradle.kts index 34c7b730b..a8ca578b4 100644 --- a/kmath-noa/build.gradle.kts +++ b/kmath-noa/build.gradle.kts @@ -27,23 +27,32 @@ val cppSources = projectDir.resolve("src/main/cpp") val cudaHome: String? = System.getenv("CUDA_HOME") val cudaDefault = file("/usr/local/cuda").exists() val cudaFound = cudaHome?.isNotEmpty() ?: false or cudaDefault - -val cmakeArchive = "cmake-3.20.5-Linux-x86_64" +val cmakeArchive = "cmake-3.20.5-linux-x86_64" val torchArchive = "libtorch" -val cmakeCmd = "$thirdPartyDir/$cmakeArchive/bin/cmake" -val ninjaCmd = "$thirdPartyDir/ninja" +val cmakeCmd = "$thirdPartyDir/cmake/$cmakeArchive/bin/cmake" +val ninjaCmd = "$thirdPartyDir/ninja/ninja" + +val generateJNIHeader by tasks.registering { + println(cmakeCmd) + doLast { + exec { + workingDir(projectDir.resolve("src/main/java/space/kscience/kmath/noa")) + commandLine("$javaHome/bin/javac", "-h", cppSources , "JNoa.java") + } + } +} val downloadCMake by tasks.registering(Download::class) { val tarFile = "$cmakeArchive.tar.gz" src("https://github.com/Kitware/CMake/releases/download/v3.20.5/$tarFile") - dest(File(thirdPartyDir, tarFile)) + dest(File("$thirdPartyDir/cmake", tarFile)) overwrite(false) } val downloadNinja by tasks.registering(Download::class) { src("https://github.com/ninja-build/ninja/releases/download/v1.10.2/ninja-linux.zip") - dest(File(thirdPartyDir, "ninja-linux.zip")) + dest(File("$thirdPartyDir/ninja", "ninja-linux.zip")) overwrite(false) } @@ -53,24 +62,72 @@ val downloadTorch by tasks.registering(Download::class) { val cpuUrl = "https://download.pytorch.org/libtorch/cpu/${torchVersion}cpu.zip" val url = if (cudaFound) cudaUrl else cpuUrl src(url) - dest(File(thirdPartyDir, "$torchArchive.zip")) + dest(File("$thirdPartyDir/torch", "$torchArchive.zip")) overwrite(false) } val extractCMake by tasks.registering(Copy::class) { dependsOn(downloadCMake) from(tarTree(resources.gzip(downloadCMake.get().dest))) - into(thirdPartyDir) + into("$thirdPartyDir/cmake") } val extractNinja by tasks.registering(Copy::class) { dependsOn(downloadNinja) from(zipTree(downloadNinja.get().dest)) - into(thirdPartyDir) + into("$thirdPartyDir/ninja") } val extractTorch by tasks.registering(Copy::class) { dependsOn(downloadTorch) from(zipTree(downloadTorch.get().dest)) - into(thirdPartyDir) -} \ No newline at end of file + into("$thirdPartyDir/torch") +} + +val configureCpp by tasks.registering { + dependsOn(extractCMake) + dependsOn(extractNinja) + dependsOn(extractTorch) + onlyIf { !file(cppBuildDir).exists() } + doLast { + exec { + workingDir(thirdPartyDir) + commandLine("mkdir", "-p", cppBuildDir) + } + exec { + workingDir(cppBuildDir) + commandLine( + cmakeCmd, + cppSources, + "-GNinja", + "-DCMAKE_MAKE_PROGRAM=$ninjaCmd", + "-DCMAKE_PREFIX_PATH=$thirdPartyDir/torch/$torchArchive", + "-DJAVA_HOME=$javaHome", + "-DCMAKE_BUILD_TYPE=Release", + "-DBUILD_NOA_CUDA=${if(!cudaFound) "ON" else "OFF"}" + ) + } + } +} + +val cleanCppBuild by tasks.registering { + onlyIf { file(cppBuildDir).exists() } + doLast { + exec { + workingDir(thirdPartyDir) + commandLine("rm", "-rf", cppBuildDir) + } + } +} + +val buildCpp by tasks.registering { + dependsOn(configureCpp) + doLast { + exec { + workingDir(cppBuildDir) + commandLine(cmakeCmd, "--build", ".", "--config", "Release") + } + } +} + +tasks["compileJava"].dependsOn(buildCpp) diff --git a/kmath-noa/src/main/cpp/CMakeLists.txt b/kmath-noa/src/main/cpp/CMakeLists.txt new file mode 100644 index 000000000..a6a70abfe --- /dev/null +++ b/kmath-noa/src/main/cpp/CMakeLists.txt @@ -0,0 +1,37 @@ +cmake_minimum_required(VERSION 3.12) + +project(JNOA LANGUAGES C CXX) + +# Require C++17 +set(CMAKE_CXX_STANDARD 17) + +find_package(Torch REQUIRED) +find_package(JNI REQUIRED) + +include(FetchContent) + +if(NOT TARGET noa) + FetchContent_Declare( + noa + GIT_REPOSITORY https://github.com/grinisrit/noa.git + GIT_TAG kmath) + + FetchContent_GetProperties(noa) + + if(NOT noa_POPULATED) + FetchContent_Populate(noa) + + set(INSTALL_NOA OFF CACHE BOOL "") + set(BUILD_NOA_TESTS OFF CACHE BOOL "") + set(BUILD_NOA_BENCHMARKS OFF CACHE BOOL "") + + add_subdirectory( + ${noa_SOURCE_DIR} + ${noa_BINARY_DIR}) + endif() +endif() + +add_library(jnoa SHARED jnoa.cc) +target_include_directories(jnoa PRIVATE ${JNI_INCLUDE_DIRS}) +target_link_libraries(jnoa PRIVATE torch NOA::NOA) +target_compile_options(jnoa PRIVATE -Wall -Wextra -Wpedantic -O3 -fPIC) diff --git a/kmath-noa/src/main/cpp/jnoa.cc b/kmath-noa/src/main/cpp/jnoa.cc new file mode 100644 index 000000000..ba446b955 --- /dev/null +++ b/kmath-noa/src/main/cpp/jnoa.cc @@ -0,0 +1,17 @@ +/* + * Copyright 2018-2021 KMath contributors. + * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. + */ + + +#include +#include + + +#include "space_kscience_kmath_noa_JNoa.h" + + +JNIEXPORT jboolean JNICALL Java_space_kscience_kmath_noa_JNoa_cudaIsAvailable(JNIEnv *, jclass) +{ + return torch::cuda::is_available(); +} \ No newline at end of file diff --git a/kmath-noa/src/main/cpp/space_kscience_kmath_noa_JNoa.h b/kmath-noa/src/main/cpp/space_kscience_kmath_noa_JNoa.h new file mode 100644 index 000000000..5554402ad --- /dev/null +++ b/kmath-noa/src/main/cpp/space_kscience_kmath_noa_JNoa.h @@ -0,0 +1,21 @@ +/* DO NOT EDIT THIS FILE - it is machine generated */ +#include +/* Header for class space_kscience_kmath_noa_JNoa */ + +#ifndef _Included_space_kscience_kmath_noa_JNoa +#define _Included_space_kscience_kmath_noa_JNoa +#ifdef __cplusplus +extern "C" { +#endif +/* + * Class: space_kscience_kmath_noa_JNoa + * Method: cudaIsAvailable + * Signature: ()Z + */ +JNIEXPORT jboolean JNICALL Java_space_kscience_kmath_noa_JNoa_cudaIsAvailable + (JNIEnv *, jclass); + +#ifdef __cplusplus +} +#endif +#endif diff --git a/kmath-noa/src/main/java/space/kscience/kmath/noa/JNoa.java b/kmath-noa/src/main/java/space/kscience/kmath/noa/JNoa.java new file mode 100644 index 000000000..ef8941848 --- /dev/null +++ b/kmath-noa/src/main/java/space/kscience/kmath/noa/JNoa.java @@ -0,0 +1,24 @@ +/* + * Copyright 2018-2021 KMath contributors. + * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. + */ + +package space.kscience.kmath.noa; + +public class JNoa { + + static { + String jNoaPath = System.getProperty("user.home") + + "/.konan/third-party/kmath-noa-0.3.0-dev-14/cpp-build/libjnoa.so"; + + try { + System.load(jNoaPath); + } catch (UnsatisfiedLinkError e) { + System.err.println("Failed to load native NOA library from:\n" + + jNoaPath +"\n" + e); + System.exit(1); + } + } + + public static native boolean cudaIsAvailable(); +} diff --git a/kmath-noa/src/main/kotlin/space/kscience/kmath/noa/utils.kt b/kmath-noa/src/main/kotlin/space/kscience/kmath/noa/utils.kt new file mode 100644 index 000000000..95cebab6b --- /dev/null +++ b/kmath-noa/src/main/kotlin/space/kscience/kmath/noa/utils.kt @@ -0,0 +1,10 @@ +/* + * Copyright 2018-2021 KMath contributors. + * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. + */ + +package space.kscience.kmath.noa + +public fun cudaAvailable(): Boolean { + return JNoa.cudaIsAvailable() +} \ No newline at end of file diff --git a/kmath-noa/src/test/kotlin/space/kscience/kmath/noa/TestUtils.kt b/kmath-noa/src/test/kotlin/space/kscience/kmath/noa/TestUtils.kt new file mode 100644 index 000000000..a3bbaa610 --- /dev/null +++ b/kmath-noa/src/test/kotlin/space/kscience/kmath/noa/TestUtils.kt @@ -0,0 +1,16 @@ +/* + * Copyright 2018-2021 KMath contributors. + * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. + */ + +package space.kscience.kmath.noa + +import kotlin.test.Test + +class TestUtils { + + @Test + fun checkCuda() { + println(cudaAvailable()) + } +} \ No newline at end of file