Skip to content

Commit

Permalink
[JNI] Add setKernelPinnedCopyThreshold and setPinnedAllocationThresho…
Browse files Browse the repository at this point in the history
…ld (#16288)

In 24.08 two new cuDF methods are being added, and the second method is still in flight (see: #16206):

```
cudf::set_kernel_pinned_copy_threshold
cudf::set_allocate_host_as_pinned_threshold
```

We'd like to expose these methods in our JNI layer. I created a Cudf.java with the two static methods, and put the definitions in CudfJni.cpp. 

Marked as draft since I need #16206 to merge, and we are still testing it.

Authors:
  - Alessandro Bellina (https://github.com/abellina)
  - Nghia Truong (https://github.com/ttnghia)

Approvers:
  - Robert (Bobby) Evans (https://github.com/revans2)
  - Jason Lowe (https://github.com/jlowe)

URL: #16288
  • Loading branch information
abellina authored Jul 24, 2024
1 parent 7191b74 commit 8fcf72a
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 0 deletions.
36 changes: 36 additions & 0 deletions java/src/main/java/ai/rapids/cudf/Cudf.java
Original file line number Diff line number Diff line change
@@ -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);
}
25 changes: 25 additions & 0 deletions java/src/main/native/src/CudfJni.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

#include <cudf/copying.hpp>
#include <cudf/utilities/default_stream.hpp>
#include <cudf/utilities/pinned_memory.hpp>

#include <sstream>

Expand Down Expand Up @@ -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<std::size_t>(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<std::size_t>(jthreshold);
cudf::set_allocate_host_as_pinned_threshold(threshold);
}
CATCH_STD(env, )
}

} // extern "C"

0 comments on commit 8fcf72a

Please sign in to comment.