diff --git a/java/src/main/java/ai/rapids/cudf/Cudf.java b/java/src/main/java/ai/rapids/cudf/Cudf.java new file mode 100644 index 00000000000..d09e2f87ed4 --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/Cudf.java @@ -0,0 +1,36 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package ai.rapids.cudf; + +public class Cudf { + + static { + NativeDepsLoader.loadNativeDeps(); + } + + /** + * cuDF copies that are smaller than the threshold will use a kernel to copy, instead + * of cudaMemcpyAsync. + */ + public static native void setKernelPinnedCopyThreshold(long kernelPinnedCopyThreshold); + + /** + * cudf allocations that are smaller than the threshold will use the pinned host + * memory resource. + */ + public static native void setPinnedAllocationThreshold(long pinnedAllocationThreshold); +} diff --git a/java/src/main/native/src/CudfJni.cpp b/java/src/main/native/src/CudfJni.cpp index 698a8f6ff02..2860dc2e4b2 100644 --- a/java/src/main/native/src/CudfJni.cpp +++ b/java/src/main/native/src/CudfJni.cpp @@ -18,6 +18,7 @@ #include #include +#include #include @@ -201,4 +202,28 @@ JNIEXPORT jboolean JNICALL Java_ai_rapids_cudf_Cuda_isPtdsEnabled(JNIEnv* env, j return cudf::jni::is_ptds_enabled; } +JNIEXPORT void JNICALL Java_ai_rapids_cudf_Cudf_setKernelPinnedCopyThreshold(JNIEnv* env, + jclass clazz, + jlong jthreshold) +{ + try { + cudf::jni::auto_set_device(env); + auto threshold = static_cast(jthreshold); + cudf::set_kernel_pinned_copy_threshold(threshold); + } + CATCH_STD(env, ) +} + +JNIEXPORT void JNICALL Java_ai_rapids_cudf_Cudf_setPinnedAllocationThreshold(JNIEnv* env, + jclass clazz, + jlong jthreshold) +{ + try { + cudf::jni::auto_set_device(env); + auto threshold = static_cast(jthreshold); + cudf::set_allocate_host_as_pinned_threshold(threshold); + } + CATCH_STD(env, ) +} + } // extern "C"