dmlc / dmlc/xgboost

Introduce typed kernel registration for device-specific implementations

Open
#12,409 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
28.8k
Forks
8.9k
Avg merge
1d 12h
Merged PRs (30d)
54

Description

## Motivation

XGBoost has a number of logical operations with separate CPU and CUDA implementations, and in some cases a SYCL implementation supplied by the plugin. Dispatch and compilation currently use a mixture of:

- Explicit `IsCUDA()` and `IsSycl()` branches.
- CPU/CUDA implementation namespaces.
- `.cc` files textually including `.cu` files.
- CMake removing core sources and recompiling them within the SYCL plugin.
- Device-specific registry tags for otherwise identical objectives or metrics.

For example, some objective `.cc` files include their corresponding `.cu` implementation. During a SYCL build, CMake moves those objective sources into the SYCL target and compiles the entire objective with `icpx`.

This couples logical operations to compiler-specific behavior and requires plugins either to modify dispatch in XGBoost proper or take ownership of core translation units.

Introduce a typed kernel registry where:

- A kernel identifies a logical operation and its function signature.
- CPU, CUDA, and plugin implementations register variants.
- The caller dispatches according to the configured device.
- Plugin-specific code remains entirely within the plugin.
- Kernels can optionally fall back to a CPU implementation.

## Proposed design

A kernel is a tag containing its function signature:

```cpp
struct HingeGradientKernel {
using Signature = void(
Context const*,
HostDeviceVector const& predictions,
MetaInfo const& info,
bst_target_t n_targets,
linalg::Matrix* output);
};

struct HingePredTransformKernel {
using Signature =
void(Context const*, HostDeviceVector* predictions);
};
```

Implementations remain ordinary functions:

```cpp
namespace cpu_impl {
void HingeGradient(/* HingeGradientKernel signature */);
}

namespace cuda_impl {
void HingeGradient(/* HingeGradientKernel signature */);
}
```

Each implementation registers itself:

```cpp
XGBOOST_REGISTER_KERNEL(
HingeGradientKernel,
IsCPU,
cpu_impl::HingeGradient);
```

```cpp
XGBOOST_REGISTER_KERNEL(
HingeGradientKernel,
IsCUDA,
cuda_impl::HingeGradient);
```

An external plugin can register another implementation:

```cpp
XGBOOST_REGISTER_KERNEL(
HingeGradientKernel,
IsSycl,
sycl_impl::HingeGradient);
```

`IsSycl` and all SYCL headers, queues, and kernels remain within the plugin.

Invocation is immediate and typed:

```cpp
DispatchKernel(
ctx_, predictions, info, Targets(info), output);
```

The objective retains configuration, validation, and other control logic:

```cpp
void GetGradient(
HostDeviceVector const& predictions,
MetaInfo const& info,
std::int32_t /*iter*/,
linalg::Matrix* output) override {
CheckInitInputs(info);
CHECK_EQ(info.labels.Size(), predictions.Size());

DispatchKernel(
ctx_, predictions, info, Targets(info), output);
}
```

## No generic kernel context

This proposal does not introduce an interface such as:

```cpp
void Run(KernelContext*);
```

XGBoost already has strongly typed data structures and eager execution. `Context` should continue to describe execution resources, while inputs and outputs remain explicit function arguments.

The useful abstraction is registered implementations of a logical kernel, not task argument packing or a task runtime.

## Device selection and fallback

A possible dispatch policy is:

1. Select an implementation supporting the requested device.
2. If none exists and the kernel permits fallback, select its CPU implementation.
3. Otherwise report that the kernel does not support the device.

The dispatcher does not need an explicit SYCL branch. The same mechanism can support other external device implementations.

Fallback policy may need to be declared per kernel because some operations can safely synchronize to the host while others should reject unsupported devices.

## Registration and linking

Registration must work with:

- Static and shared XGBoost builds.
- CUDA translation units.
- Plugin object libraries.
- Builds where optional backends are absent.

Duplicate registrations for the same kernel and device should be rejected.

Registration translation units must also be retained during static linking. This could build on the existing registry/link-tag mechanism or use explicit object inclusion.

Registered CUDA and SYCL implementations are host-callable wrappers. Those wrappers prepare device data and launch their backend kernels.

## Initial example: binary hinge

Hinge provides a small proof of concept with two kernels:

```text
HingeGradientKernel
HingePredTransformKernel
```

Proposed source layout:

```text
src/objective/hinge.h
Shared hinge calculations
Kernel signatures

src/objective/hinge.cc
Objective class and registration
CPU kernel implementations and registration

src/objective/hinge.cu
CUDA kernel implementations and registration

plugin/sycl/objective/hinge.cc
Optional SYCL implementations and registration
```

This removes the need to compile the entire objective with every backend compiler.

## Other potential users

If the hinge prototype is successful, candidates include:

- Other objective gradient and prediction-transform operations.
- Initial-score/stump fitting.
- Statistical reductions such as mean and median.
- Multiclass and survival metrics.
- SHAP value and interaction calculations.
- Small multi-target tree operations.

Each operation defines its own signature. Specialized and rowwise objectives do not need to conform to an elementwise interface.

## Non-goals

This proposal does not attempt to:

- Introduce deferred execution or a task graph.
- Add a scheduler.
- Replace `Context`.
- Introduce a generic argument container.
- Refactor `HostDeviceVector` storage or coherence.
- Define a stable plugin ABI.
- Replace existing predictor or updater registries.
- Runtime-dispatch arbitrary templated lambdas.
- Require objectives to use a common elementwise abstraction.

Low-level templated utilities such as `ElementWiseKernel` remain compile-time mechanisms. The registered kernel represents the higher-level semantic operation that owns those device-specific loops.

## Open questions

- Should every kernel require a CPU implementation?
- Should fallback be automatic or declared per kernel?
- Should variants use exact device kinds or matching predicates?
- Should the registry build on the existing DMLC registry?
- How should backend-specific caches be associated with variants?
- Should kernel declarations be public or internal?
- How should duplicate registrations and missing variants be diagnosed?

## Suggested implementation sequence

1. Add kernel registration and dispatch.
2. Convert hinge gradient and prediction transform.
3. Verify CPU-only, CUDA, static, and shared builds.
4. Verify SYCL CPU fallback.
5. Add a plugin-provided SYCL hinge implementation.
6. Convert `FitStump` as a second kernel with a different signature.
7. Evaluate the abstraction against additional objectives and reductions.

Contributor guide

No contributing guide indexed for this repository

Research direction

Read the proposed kernel layout in src/objective/hinge.h, src/objective/hinge.cc, src/objective/hinge.cu, and plugin/sycl/objective/hinge.cc, then start with the registration and dispatch design. Use the hinge gradient and prediction-transform kernels as the initial scope. Done includes CPU-only, CUDA, static, shared, SYCL fallback, and plugin-provided SYCL build verification.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
backend, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.