Simple tests.

Author: trivialfis

src/common/linalg_op.h

@@ -5,6 +5,9 @@
  *
  * Client code can use utilities like @ref ElementWiseKernel by including this file in the
  * right translation unit. For CUDA-compatible kernels, include this header in a .cu TU.
+ *
+ * Be aware of potential violation of the one definition rule (ODR). The dispatching
+ * functions should never be used in an inline function without a system tag.
  */
 #ifndef XGBOOST_COMMON_LINALG_OP_H_
 #define XGBOOST_COMMON_LINALG_OP_H_

tests/cpp/common/test_linalg.cc

@@ -11,6 +11,7 @@
 #include <vector>   // for vector
 
 #include "../../../src/common/linalg_op.h"
+#include "test_linalg.h"  // for TestLinalgDispatch
 
 namespace xgboost::linalg {
 namespace {
@@ -410,4 +411,9 @@ TEST(Linalg, IO) {
     check(loaded);
   }
 }
+
+TEST(Linalg, CpuDispatch) {
+  Context ctx;
+  TestLinalgDispatch(&ctx, [](auto v) { return v + 1; });
+}
 }  // namespace xgboost::linalg

tests/cpp/common/test_linalg.cu

@@ -10,6 +10,7 @@
 #include "../../../src/common/linalg_op.h"
 #include "../../../src/common/optional_weight.h"  // for MakeOptionalWeights
 #include "../helpers.h"
+#include "test_linalg.h"     // for TestLinalgDispatch
 #include "thrust/random.h"   // for default_random_engine
 #include "thrust/shuffle.h"  // for shuffle
 #include "xgboost/context.h"
@@ -28,8 +29,7 @@ void TestElementWiseKernel() {
     // GPU view
     auto t = l.View(device).Slice(linalg::All(), 1, linalg::All());
     ASSERT_FALSE(t.CContiguous());
-    cuda_impl::TransformIdxKernel(&ctx, t,
-                                          [] XGBOOST_DEVICE(std::size_t i, float) { return i; });
+    cuda_impl::TransformIdxKernel(&ctx, t, [] XGBOOST_DEVICE(std::size_t i, float) { return i; });
     // CPU view
     t = l.View(DeviceOrd::CPU()).Slice(linalg::All(), 1, linalg::All());
     std::size_t k = 0;
@@ -56,8 +56,7 @@ void TestElementWiseKernel() {
      * Contiguous
      */
     auto t = l.View(device);
-    cuda_impl::TransformIdxKernel(&ctx, t,
-                                          [] XGBOOST_DEVICE(size_t i, float) { return i; });
+    cuda_impl::TransformIdxKernel(&ctx, t, [] XGBOOST_DEVICE(size_t i, float) { return i; });
     ASSERT_TRUE(t.CContiguous());
     // CPU view
     t = l.View(DeviceOrd::CPU());
@@ -151,4 +150,11 @@ TEST(Linalg, SmallHistogram) {
     ASSERT_EQ(h_bins[i], cnt);
   }
 }
+namespace {
+void TestGpuDispatch() {
+  auto ctx = MakeCUDACtx(0);
+  TestLinalgDispatch(&ctx, [] XGBOOST_DEVICE(double v) { return v + 1; });
+}
+}  // namespace
+TEST(Linalg, GpuDispatch) { TestGpuDispatch(); }
 }  // namespace xgboost::linalg

tests/cpp/common/test_linalg.h

@@ -0,0 +1,28 @@
+/**
+ * Copyright 2025, XGBoost Contributors
+ */
+#pragma once
+
+#include <gtest/gtest.h>
+#include <xgboost/context.h>
+#include <xgboost/linalg.h>  // for Vector
+
+#include <numeric>  // for iota
+#include <vector>   // for vector
+
+#include "../../../src/common/linalg_op.h"
+
+namespace xgboost::linalg {
+template <typename Fn>
+void TestLinalgDispatch(Context const* ctx, Fn&& fn) {
+  std::vector<double> data(128, 0);
+  std::iota(data.begin(), data.end(), 0.0);
+  Vector<double> vec(data.begin(), data.end(), {data.size()}, DeviceOrd::CPU());
+
+  TransformKernel(ctx, vec.View(ctx->Device()), [=] XGBOOST_DEVICE(double v) { return fn(v); });
+  auto h_v = vec.HostView();
+  for (std::size_t i = 0; i < h_v.Size(); ++i) {
+    ASSERT_EQ(h_v(i), fn(i));
+  }
+}
+}  // namespace xgboost::linalg