Unittest.

tests/cpp/common/test_linalg.cu

@@ -0,0 +1,62 @@
+/*!
+ * Copyright 2021 by XGBoost Contributors
+ */
+#include <gtest/gtest.h>
+
+#include "../../../src/common/linalg_op.cuh"
+#include "xgboost/generic_parameters.h"
+#include "xgboost/linalg.h"
+
+namespace xgboost {
+namespace linalg {
+namespace {
+void TestElementWiseKernel() {
+  Tensor<float, 3> l{{2, 3, 4}, 0};
+  {
+    /**
+     * Non-contiguous
+     */
+    // GPU view
+    auto t = l.View(0).Slice(linalg::All(), 1, linalg::All());
+    ASSERT_FALSE(t.Contiguous());
+    ElementWiseKernelDevice(t, [] __device__(size_t i, float) { return i; });
+    // CPU view
+    t = l.View(GenericParameter::kCpuId).Slice(linalg::All(), 1, linalg::All());
+    size_t k = 0;
+    for (size_t i = 0; i < l.Shape(0); ++i) {
+      for (size_t j = 0; j < l.Shape(2); ++j) {
+        ASSERT_EQ(k++, t(i, j));
+      }
+    }
+
+    t = l.View(0).Slice(linalg::All(), 1, linalg::All());
+    ElementWiseKernelDevice(t, [] __device__(size_t i, float v) {
+      SPAN_CHECK(v == i);
+      return v;
+    });
+  }
+
+  {
+    /**
+     * Contiguous
+     */
+    auto t = l.View(0);
+    ElementWiseKernelDevice(t, [] __device__(size_t i, float) { return i; });
+    ASSERT_TRUE(t.Contiguous());
+    // CPU view
+    t = l.View(GenericParameter::kCpuId);
+
+    size_t ind = 0;
+    for (size_t i = 0; i < l.Shape(0); ++i) {
+      for (size_t j = 0; j < l.Shape(1); ++j) {
+        for (size_t k = 0; k < l.Shape(2); ++k) {
+          ASSERT_EQ(ind++, t(i, j, k));
+        }
+      }
+    }
+  }
+}
+}  // anonymous namespace
+TEST(Linalg, GPUElementWise) { TestElementWiseKernel(); }
+}  // namespace linalg
+}  // namespace xgboost