xgboost

linalg_op.h
Go to the documentation of this file.
1
4 #ifndef XGBOOST_COMMON_LINALG_OP_H_
5 #define XGBOOST_COMMON_LINALG_OP_H_
6 #include <type_traits>
7
8 #include "common.h"
9 #include "threading_utils.h"
10 #include "xgboost/generic_parameters.h"
11 #include "xgboost/linalg.h"
12
13 namespace xgboost {
14 namespace linalg {
15 template <typename T, int32_t D, typename Fn>
16 void ElementWiseTransformHost(linalg::TensorView<T, D> t, int32_t n_threads, Fn&& fn) {
17 if (t.Contiguous()) {
18 auto ptr = t.Values().data();
19 common::ParallelFor(t.Size(), n_threads, [&](size_t i) { ptr[i] = fn(i, ptr[i]); });
20 } else {
21 common::ParallelFor(t.Size(), n_threads, [&](size_t i) {
22 auto& v = detail::Apply(t, linalg::UnravelIndex(i, t.Shape()));
23 v = fn(i, v);
24 });
25 }
26 }
27
28 template <typename T, int32_t D, typename Fn>
29 void ElementWiseKernelHost(linalg::TensorView<T, D> t, int32_t n_threads, Fn&& fn) {
30 static_assert(std::is_void<std::result_of_t<Fn(size_t, T&)>>::value,
31 "For function with return, use transform instead.");
32 if (t.Contiguous()) {
33 auto ptr = t.Values().data();
34 common::ParallelFor(t.Size(), n_threads, [&](size_t i) { fn(i, ptr[i]); });
35 } else {
36 common::ParallelFor(t.Size(), n_threads, [&](size_t i) {
37 auto& v = detail::Apply(t, linalg::UnravelIndex(i, t.Shape()));
38 fn(i, v);
39 });
40 }
41 }
42
43 #if !defined(XGBOOST_USE_CUDA)
44 template <typename T, int32_t D, typename Fn>
45 void ElementWiseKernelDevice(linalg::TensorView<T, D> t, Fn&& fn, void* s = nullptr) {
46 common::AssertGPUSupport();
47 }
48
49 template <typename T, int32_t D, typename Fn>
50 void ElementWiseTransformDevice(linalg::TensorView<T, D> t, Fn&& fn, void* s = nullptr) {
51 common::AssertGPUSupport();
52 }
53
54 template <typename T, int32_t D, typename Fn>
55 void ElementWiseKernel(GenericParameter const* ctx, linalg::TensorView<T, D> t, Fn&& fn) {
56 if (!ctx->IsCPU()) {
57 common::AssertGPUSupport();
58 }
59 ElementWiseKernelHost(t, ctx->Threads(), fn);
60 }
61 #endif // !defined(XGBOOST_USE_CUDA)
62 } // namespace linalg
63 } // namespace xgboost
64 #endif // XGBOOST_COMMON_LINALG_OP_H_
xgboost::linalg::TensorView
A tensor view with static type and dimension. It implements indexing and slicing.
Definition: linalg.h:244
xgboost::linalg::TensorView::Size
LINALG_HD size_t Size() const
Number of items in the tensor.
Definition: linalg.h:467
xgboost::linalg::ElementWiseKernelDevice
void ElementWiseKernelDevice(linalg::TensorView< T, D > t, Fn &&fn, void *s=nullptr)
Definition: linalg_op.h:45
xgboost::linalg::ElementWiseKernel
void ElementWiseKernel(GenericParameter const *ctx, linalg::TensorView< T, D > t, Fn &&fn)
Definition: linalg_op.h:55
xgboost::GenericParameter
Definition: generic_parameters.h:15
xgboost::linalg::ElementWiseTransformHost
void ElementWiseTransformHost(linalg::TensorView< T, D > t, int32_t n_threads, Fn &&fn)
Definition: linalg_op.h:16
xgboost::linalg::TensorView::Values
LINALG_HD auto Values() const -> decltype(data_) const &
Obtain a reference to the raw data.
Definition: linalg.h:497
xgboost::linalg::ElementWiseKernelHost
void ElementWiseKernelHost(linalg::TensorView< T, D > t, int32_t n_threads, Fn &&fn)
Definition: linalg_op.h:29
xgboost::common::ParallelFor
void ParallelFor(Index size, int32_t n_threads, Sched sched, Func fn)
Definition: threading_utils.h:170
xgboost::common::AssertGPUSupport
void AssertGPUSupport()
Definition: common.h:237
linalg.h
Linear algebra related utilities.
xgboost::linalg::TensorView::Contiguous
LINALG_HD bool Contiguous() const
Whether this is a contiguous array, both C and F contiguous returns true.
Definition: linalg.h:471
xgboost::linalg::ElementWiseTransformDevice
void ElementWiseTransformDevice(linalg::TensorView< T, D > t, Fn &&fn, void *s=nullptr)
Definition: linalg_op.h:50
common.h
Common utilities.
generic_parameters.h
threading_utils.h
xgboost::GenericParameter::Threads
int32_t Threads() const
xgboost::GenericParameter::IsCPU
bool IsCPU() const
Definition: generic_parameters.h:52
xgboost
namespace of xgboost
Definition: base.h:110