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"
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) {
47 }
48 
49 template <typename T, int32_t D, typename Fn>
50 void ElementWiseTransformDevice(linalg::TensorView<T, D> t, Fn&& fn, void* s = nullptr) {
52 }
53 
54 template <typename T, int32_t D, typename Fn>
56  if (!ctx->IsCPU()) {
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