Tensor: Fix output buffer initialization on zero-size contraction dimension libeigen/eigen!2898
diff --git a/unsupported/Eigen/src/Tensor/TensorContraction.h b/unsupported/Eigen/src/Tensor/TensorContraction.h index c67c8a7..f5d86d9 100644 --- a/unsupported/Eigen/src/Tensor/TensorContraction.h +++ b/unsupported/Eigen/src/Tensor/TensorContraction.h
@@ -684,6 +684,17 @@ template <bool lhs_inner_dim_contiguous, bool rhs_inner_dim_contiguous, bool rhs_inner_dim_reordered, int Alignment> void evalProductSequential(Scalar* buffer) const { + if (this->m_i_size == 0 || this->m_j_size == 0) { + return; + } + if (this->m_k_size == 0) { + // Contraction over dimension of size zero results in an all-zero output. + this->m_device.fill(buffer, buffer + this->m_i_size * this->m_j_size, Scalar(0)); + using OutputMapper = internal::blas_data_mapper<Scalar, Index, ColMajor>; + this->m_output_kernel(OutputMapper(buffer, this->m_i_size), this->m_tensor_contraction_params, + static_cast<Index>(0), static_cast<Index>(0), this->m_i_size, this->m_j_size); + return; + } // Gemv-shape contractions (output is a vector) get a direct GEMV kernel // call when the operand layout admits one. There are four shapes: // @@ -856,6 +867,17 @@ // columns in right side const Index n = this->m_j_size; + if (m == 0 || n == 0) return; + if (k_slice == 0) { + this->m_device.fill(buffer, buffer + m * n, Scalar(0)); + if (use_output_kernel) { + using OutputMapper = internal::blas_data_mapper<Scalar, Index, ColMajor>; + m_output_kernel(OutputMapper(buffer, m), m_tensor_contraction_params, static_cast<Index>(0), + static_cast<Index>(0), m, n); + } + return; + } + // define data mappers for Lhs and Rhs using LhsScalar = std::remove_const_t<typename EvalLeftArgType::Scalar>; using RhsScalar = std::remove_const_t<typename EvalRightArgType::Scalar>;
diff --git a/unsupported/Eigen/src/Tensor/TensorContractionGpu.h b/unsupported/Eigen/src/Tensor/TensorContractionGpu.h index 048bed8..8e4fac4 100644 --- a/unsupported/Eigen/src/Tensor/TensorContractionGpu.h +++ b/unsupported/Eigen/src/Tensor/TensorContractionGpu.h
@@ -1326,6 +1326,11 @@ } } + template <int Alignment> + void evalProduct(Scalar* buffer) const { + evalTo(buffer); + } + template <typename LhsScalar, typename RhsScalar, typename Index, typename LhsMapper, typename RhsMapper, typename OutputMapper, bool UseNaiveKernel> struct LaunchKernelsImpl; @@ -1389,16 +1394,19 @@ void evalTyped(Scalar* buffer) const { // columns in left side, rows in right side const Index k = this->m_k_size; - EIGEN_UNUSED_VARIABLE(k); // rows in left side const Index m = this->m_i_size; // columns in right side const Index n = this->m_j_size; + if (m == 0 || n == 0) return; + // zero out the result buffer (which must be of size at least m * n * sizeof(Scalar)) this->m_device.fill(buffer, buffer + m * n, Scalar(0)); + if (k == 0) return; + typedef internal::TensorContractionInputMapper<LhsScalar, Index, internal::Lhs, LeftEvaluator, left_nocontract_t, contract_t, 4, lhs_inner_dim_contiguous, false, Unaligned> LhsMapper;
diff --git a/unsupported/Eigen/src/Tensor/TensorContractionThreadPool.h b/unsupported/Eigen/src/Tensor/TensorContractionThreadPool.h index 6fdee7d..96da0ca 100644 --- a/unsupported/Eigen/src/Tensor/TensorContractionThreadPool.h +++ b/unsupported/Eigen/src/Tensor/TensorContractionThreadPool.h
@@ -105,7 +105,19 @@ const Index m = this->m_i_size; const Index n = this->m_j_size; const Index k = this->m_k_size; - if (m == 0 || n == 0 || k == 0) return; + if (m == 0 || n == 0) { + EIGEN_IF_CONSTEXPR (!IsEvalInSyncMode) done(); + return; + } + if (k == 0) { + internal::tensor_contraction_dispatch( + [&](auto lhs_c, auto rhs_c, auto rhs_r) { + this->template evalProductSequential<lhs_c(), rhs_c(), rhs_r(), Unaligned>(buffer); + }, + this->m_lhs_inner_dim_contiguous, this->m_rhs_inner_dim_contiguous, this->m_rhs_inner_dim_reordered); + EIGEN_IF_CONSTEXPR (!IsEvalInSyncMode) done(); + return; + } // Compute a set of algorithm parameters: // - kernel block sizes (bm, bn, bk)
diff --git a/unsupported/test/tensor_contract_gpu.cu b/unsupported/test/tensor_contract_gpu.cu index 07b0eeb..1db50b1 100644 --- a/unsupported/test/tensor_contract_gpu.cu +++ b/unsupported/test/tensor_contract_gpu.cu
@@ -75,6 +75,52 @@ gpuFree((void*)d_t_result); } +template <int DataLayout, typename Scalar = float> +void test_gpu_contraction_zero_k(int m_size, int n_size) { + const int k_size = 0; + Tensor<Scalar, 2, DataLayout> t_left(m_size, k_size); + Tensor<Scalar, 2, DataLayout> t_right(k_size, n_size); + Tensor<Scalar, 2, DataLayout> t_result(m_size, n_size); + Tensor<Scalar, 2, DataLayout> t_result_gpu(m_size, n_size); + t_result_gpu.setConstant(Scalar(123)); + Eigen::array<DimPair, 1> dims{DimPair(1, 0)}; + + std::size_t t_left_bytes = (std::max)(t_left.size(), DenseIndex(1)) * sizeof(Scalar); + std::size_t t_right_bytes = (std::max)(t_right.size(), DenseIndex(1)) * sizeof(Scalar); + std::size_t t_result_bytes = t_result.size() * sizeof(Scalar); + + Scalar* d_t_left; + Scalar* d_t_right; + Scalar* d_t_result; + + gpuMalloc((void**)(&d_t_left), t_left_bytes); + gpuMalloc((void**)(&d_t_right), t_right_bytes); + gpuMalloc((void**)(&d_t_result), t_result_bytes); + + // Prefill the device output with a nonzero value to verify that the GPU path actively zeroes the output buffer. + gpuMemcpy(d_t_result, t_result_gpu.data(), t_result_bytes, gpuMemcpyHostToDevice); + + Eigen::GpuStreamDevice stream; + Eigen::GpuDevice gpu_device(&stream); + + Eigen::TensorMap<Eigen::Tensor<Scalar, 2, DataLayout>> gpu_t_left(d_t_left, Eigen::array<int, 2>{m_size, k_size}); + Eigen::TensorMap<Eigen::Tensor<Scalar, 2, DataLayout>> gpu_t_right(d_t_right, Eigen::array<int, 2>{k_size, n_size}); + Eigen::TensorMap<Eigen::Tensor<Scalar, 2, DataLayout>> gpu_t_result(d_t_result, Eigen::array<int, 2>{m_size, n_size}); + + gpu_t_result.device(gpu_device) = gpu_t_left.contract(gpu_t_right, dims); + t_result = t_left.contract(t_right, dims); + + gpuMemcpy(t_result_gpu.data(), d_t_result, t_result_bytes, gpuMemcpyDeviceToHost); + for (DenseIndex i = 0; i < t_result.size(); i++) { + VERIFY_IS_EQUAL(t_result(i), Scalar(0)); + VERIFY_IS_EQUAL(t_result_gpu(i), Scalar(0)); + } + + gpuFree((void*)d_t_left); + gpuFree((void*)d_t_right); + gpuFree((void*)d_t_result); +} + template <int DataLayout> void test_gpu_contraction_double(int m_size, int k_size, int n_size) { Tensor<double, 2, DataLayout> t_left(m_size, k_size); @@ -256,6 +302,15 @@ CALL_SUBTEST_1(test_scalar<ColMajor>(128, 128, 128)); CALL_SUBTEST_1(test_scalar<RowMajor>(128, 128, 128)); + CALL_SUBTEST_1((test_gpu_contraction_zero_k<ColMajor>(10, 20))); + CALL_SUBTEST_1((test_gpu_contraction_zero_k<RowMajor>(10, 20))); + CALL_SUBTEST_1((test_gpu_contraction_zero_k<ColMajor>(64, 64))); + CALL_SUBTEST_1((test_gpu_contraction_zero_k<RowMajor>(64, 64))); + CALL_SUBTEST_1((test_gpu_contraction_zero_k<ColMajor, double>(10, 20))); + CALL_SUBTEST_1((test_gpu_contraction_zero_k<RowMajor, double>(10, 20))); + CALL_SUBTEST_1((test_gpu_contraction_zero_k<ColMajor, double>(64, 64))); + CALL_SUBTEST_1((test_gpu_contraction_zero_k<RowMajor, double>(64, 64))); + CALL_SUBTEST_2(test_gpu_contraction_m<ColMajor>()); CALL_SUBTEST_3(test_gpu_contraction_m<RowMajor>());
diff --git a/unsupported/test/tensor_contraction.cpp b/unsupported/test/tensor_contraction.cpp index bed0dae..a2616df 100644 --- a/unsupported/test/tensor_contraction.cpp +++ b/unsupported/test/tensor_contraction.cpp
@@ -593,6 +593,334 @@ } } +template <int DataLayout, typename Scalar> +static void test_zero_dim_contraction_2d() { + // Test {10, 0} * {0, 20} -> {10, 20} + Tensor<Scalar, 2, DataLayout> t_left(10, 0); + Tensor<Scalar, 2, DataLayout> t_right(0, 20); + + // Pre-fill destination tensor with non-zero values to verify that contraction actively zeroes output memory. + const Scalar non_zero = Scalar(123); + Tensor<Scalar, 2, DataLayout> t_result(10, 20); + t_result.setConstant(non_zero); + + Eigen::array<DimPair, 1> dims = {{DimPair(1, 0)}}; + t_result = t_left.contract(t_right, dims); + + VERIFY_IS_EQUAL(t_result.dimension(0), 10); + VERIFY_IS_EQUAL(t_result.dimension(1), 20); + VERIFY_IS_EQUAL(t_result.size(), 200); + + for (Index i = 0; i < t_result.dimension(0); ++i) { + for (Index j = 0; j < t_result.dimension(1); ++j) { + VERIFY_IS_EQUAL(t_result(i, j), Scalar(0)); + } + } + + // Also test via TensorEvaluator directly with evalTo into pre-filled memory. + Tensor<Scalar, 2, DataLayout> t_result_eval(10, 20); + t_result_eval.setConstant(non_zero); + typedef TensorEvaluator<decltype(t_left.contract(t_right, dims)), DefaultDevice> Evaluator; + Evaluator eval(t_left.contract(t_right, dims), DefaultDevice()); + eval.evalTo(t_result_eval.data()); + for (Index i = 0; i < t_result_eval.size(); ++i) { + VERIFY_IS_EQUAL(t_result_eval.data()[i], Scalar(0)); + } + + // Test large contraction: {64, 0} * {0, 64} -> {64, 64} + // This exercises the path where max(m, n) >= 48, which previously triggered a + // SIGFPE (division by zero) in computeProductBlockingSizes. + { + Tensor<Scalar, 2, DataLayout> t_left_large(64, 0); + Tensor<Scalar, 2, DataLayout> t_right_large(0, 64); + Tensor<Scalar, 2, DataLayout> t_result_large(64, 64); + t_result_large.setConstant(non_zero); + + t_result_large = t_left_large.contract(t_right_large, dims); + + VERIFY_IS_EQUAL(t_result_large.dimension(0), 64); + VERIFY_IS_EQUAL(t_result_large.dimension(1), 64); + VERIFY_IS_EQUAL(t_result_large.size(), 64 * 64); + + for (Index i = 0; i < t_result_large.size(); ++i) { + VERIFY_IS_EQUAL(t_result_large.data()[i], Scalar(0)); + } + } + + // 1D Vector contractions (GEMV shapes) + // Matrix {10, 0} * Vector {0} -> Vector {10} + Tensor<Scalar, 1, DataLayout> vec_right(0); + Tensor<Scalar, 1, DataLayout> vec_result(10); + vec_result.setConstant(non_zero); + vec_result = t_left.contract(vec_right, dims); + VERIFY_IS_EQUAL(vec_result.dimension(0), 10); + for (Index i = 0; i < vec_result.size(); ++i) { + VERIFY_IS_EQUAL(vec_result(i), Scalar(0)); + } + + // Vector {0} * Matrix {0, 20} -> Vector {20} + Tensor<Scalar, 1, DataLayout> vec_left(0); + Tensor<Scalar, 1, DataLayout> vec_result2(20); + vec_result2.setConstant(non_zero); + Eigen::array<DimPair, 1> dims_vm = {{DimPair(0, 0)}}; + vec_result2 = vec_left.contract(t_right, dims_vm); + VERIFY_IS_EQUAL(vec_result2.dimension(0), 20); + for (Index i = 0; i < vec_result2.size(); ++i) { + VERIFY_IS_EQUAL(vec_result2(i), Scalar(0)); + } + + // Vector {0} * Vector {0} -> Scalar {} + Tensor<Scalar, 0, DataLayout> scalar_result; + scalar_result() = non_zero; + scalar_result = vec_left.contract(vec_right, dims_vm); + VERIFY_IS_EQUAL(scalar_result(), Scalar(0)); +} + +template <int DataLayout, typename Scalar> +static void test_zero_dim_contraction_multidims() { + // Test {5, 0, 4} * {0, 3, 4} with contraction dims {1, 0} and {2, 2} -> {5, 3} + Tensor<Scalar, 3, DataLayout> t_left(5, 0, 4); + Tensor<Scalar, 3, DataLayout> t_right(0, 3, 4); + + const Scalar non_zero = Scalar(123); + Tensor<Scalar, 2, DataLayout> t_result(5, 3); + t_result.setConstant(non_zero); + + Eigen::array<DimPair, 2> dims = {{DimPair(1, 0), DimPair(2, 2)}}; + t_result = t_left.contract(t_right, dims); + + VERIFY_IS_EQUAL(t_result.dimension(0), 5); + VERIFY_IS_EQUAL(t_result.dimension(1), 3); + VERIFY_IS_EQUAL(t_result.size(), 15); + + for (Index i = 0; i < t_result.dimension(0); ++i) { + for (Index j = 0; j < t_result.dimension(1); ++j) { + VERIFY_IS_EQUAL(t_result(i, j), Scalar(0)); + } + } + + // 4D contractions: {2, 5, 0, 3} * {0, 4, 3, 2} with contraction dims {2, 0} and {3, 2} -> {2, 5, 4, 2} + Tensor<Scalar, 4, DataLayout> t_left4(2, 5, 0, 3); + Tensor<Scalar, 4, DataLayout> t_right4(0, 4, 3, 2); + Tensor<Scalar, 4, DataLayout> t_result4(2, 5, 4, 2); + t_result4.setConstant(non_zero); + + Eigen::array<DimPair, 2> dims4 = {{DimPair(2, 0), DimPair(3, 2)}}; + t_result4 = t_left4.contract(t_right4, dims4); + + VERIFY_IS_EQUAL(t_result4.dimension(0), 2); + VERIFY_IS_EQUAL(t_result4.dimension(1), 5); + VERIFY_IS_EQUAL(t_result4.dimension(2), 4); + VERIFY_IS_EQUAL(t_result4.dimension(3), 2); + VERIFY_IS_EQUAL(t_result4.size(), 80); + + for (Index i = 0; i < t_result4.size(); ++i) { + VERIFY_IS_EQUAL(t_result4.data()[i], Scalar(0)); + } +} + +template <int DataLayout, typename Scalar> +static void test_zero_dim_contraction_outer_zeros() { + Eigen::array<DimPair, 1> dims = {{DimPair(1, 0)}}; + + // Outer zero on LHS: {0, 5} * {5, 10} -> {0, 10} + { + Tensor<Scalar, 2, DataLayout> t_left(0, 5); + Tensor<Scalar, 2, DataLayout> t_right(5, 10); + t_right.setRandom(); + Tensor<Scalar, 2, DataLayout> t_result(0, 10); + t_result = t_left.contract(t_right, dims); + VERIFY_IS_EQUAL(t_result.dimension(0), 0); + VERIFY_IS_EQUAL(t_result.dimension(1), 10); + VERIFY_IS_EQUAL(t_result.size(), 0); + } + + // Outer zero on RHS: {10, 5} * {5, 0} -> {10, 0} + { + Tensor<Scalar, 2, DataLayout> t_left(10, 5); + Tensor<Scalar, 2, DataLayout> t_right(5, 0); + t_left.setRandom(); + Tensor<Scalar, 2, DataLayout> t_result(10, 0); + t_result = t_left.contract(t_right, dims); + VERIFY_IS_EQUAL(t_result.dimension(0), 10); + VERIFY_IS_EQUAL(t_result.dimension(1), 0); + VERIFY_IS_EQUAL(t_result.size(), 0); + } + + // Outer zero on both LHS and RHS: {0, 5} * {5, 0} -> {0, 0} + { + Tensor<Scalar, 2, DataLayout> t_left(0, 5); + Tensor<Scalar, 2, DataLayout> t_right(5, 0); + Tensor<Scalar, 2, DataLayout> t_result(0, 0); + t_result = t_left.contract(t_right, dims); + VERIFY_IS_EQUAL(t_result.dimension(0), 0); + VERIFY_IS_EQUAL(t_result.dimension(1), 0); + VERIFY_IS_EQUAL(t_result.size(), 0); + } + + // Both outer and inner zero: {0, 0} * {0, 10} -> {0, 10} + { + Tensor<Scalar, 2, DataLayout> t_left(0, 0); + Tensor<Scalar, 2, DataLayout> t_right(0, 10); + Tensor<Scalar, 2, DataLayout> t_result(0, 10); + t_result = t_left.contract(t_right, dims); + VERIFY_IS_EQUAL(t_result.dimension(0), 0); + VERIFY_IS_EQUAL(t_result.dimension(1), 10); + VERIFY_IS_EQUAL(t_result.size(), 0); + } + + // Both outer and inner zero: {10, 0} * {0, 0} -> {10, 0} + { + Tensor<Scalar, 2, DataLayout> t_left(10, 0); + Tensor<Scalar, 2, DataLayout> t_right(0, 0); + Tensor<Scalar, 2, DataLayout> t_result(10, 0); + t_result = t_left.contract(t_right, dims); + VERIFY_IS_EQUAL(t_result.dimension(0), 10); + VERIFY_IS_EQUAL(t_result.dimension(1), 0); + VERIFY_IS_EQUAL(t_result.size(), 0); + } + + // Both outer and inner zero: {0, 0} * {0, 0} -> {0, 0} + { + Tensor<Scalar, 2, DataLayout> t_left(0, 0); + Tensor<Scalar, 2, DataLayout> t_right(0, 0); + Tensor<Scalar, 2, DataLayout> t_result(0, 0); + t_result = t_left.contract(t_right, dims); + VERIFY_IS_EQUAL(t_result.dimension(0), 0); + VERIFY_IS_EQUAL(t_result.dimension(1), 0); + VERIFY_IS_EQUAL(t_result.size(), 0); + } + + // Multidimensional with outer zero: {0, 3, 4} * {3, 2, 4} with contraction dims {1, 0} and {2, 2} -> {0, 2} + { + Tensor<Scalar, 3, DataLayout> t_left(0, 3, 4); + Tensor<Scalar, 3, DataLayout> t_right(3, 2, 4); + t_right.setRandom(); + Eigen::array<DimPair, 2> dims_md = {{DimPair(1, 0), DimPair(2, 2)}}; + Tensor<Scalar, 2, DataLayout> t_result(0, 2); + t_result = t_left.contract(t_right, dims_md); + VERIFY_IS_EQUAL(t_result.dimension(0), 0); + VERIFY_IS_EQUAL(t_result.dimension(1), 2); + VERIFY_IS_EQUAL(t_result.size(), 0); + } +} + +template <typename Scalar> +struct AddBiasOutputKernel { + Scalar bias; + explicit AddBiasOutputKernel(Scalar b) : bias(b) {} + + template <typename Index, typename ResScalar> + EIGEN_ALWAYS_INLINE void operator()(const internal::blas_data_mapper<ResScalar, Index, ColMajor>& output_mapper, + const TensorContractionParams&, Index, Index, Index num_rows, + Index num_cols) const { + for (Index i = 0; i < num_rows; ++i) { + for (Index j = 0; j < num_cols; ++j) { + output_mapper(i, j) += bias; + } + } + } +}; + +template <int DataLayout, typename Scalar> +static void test_zero_dim_contraction_output_kernel() { + const Scalar non_zero = Scalar(123); + const Scalar bias = Scalar(42); + + // 2D contraction with bias addition on zero-sized contraction dimension: + // {10, 0} * {0, 20} -> {10, 20} + // The contraction sum is 0, so output kernel adds bias to produce `bias` everywhere. + { + Tensor<Scalar, 2, DataLayout> t_left(10, 0); + Tensor<Scalar, 2, DataLayout> t_right(0, 20); + Tensor<Scalar, 2, DataLayout> t_result(10, 20); + t_result.setConstant(non_zero); + + Eigen::array<DimPair, 1> dims = {{DimPair(1, 0)}}; + t_result = t_left.contract(t_right, dims, AddBiasOutputKernel<Scalar>(bias)); + + VERIFY_IS_EQUAL(t_result.dimension(0), 10); + VERIFY_IS_EQUAL(t_result.dimension(1), 20); + for (Index i = 0; i < t_result.dimension(0); ++i) { + for (Index j = 0; j < t_result.dimension(1); ++j) { + VERIFY_IS_EQUAL(t_result(i, j), bias); + } + } + } + + // 2D contraction with SqrtOutputKernel on zero-sized contraction dimension: + // sqrt(0) == 0 + { + Tensor<Scalar, 2, DataLayout> t_left(10, 0); + Tensor<Scalar, 2, DataLayout> t_right(0, 20); + Tensor<Scalar, 2, DataLayout> t_result(10, 20); + t_result.setConstant(non_zero); + + Eigen::array<DimPair, 1> dims = {{DimPair(1, 0)}}; + t_result = t_left.contract(t_right, dims, SqrtOutputKernel()); + + VERIFY_IS_EQUAL(t_result.dimension(0), 10); + VERIFY_IS_EQUAL(t_result.dimension(1), 20); + for (Index i = 0; i < t_result.size(); ++i) { + VERIFY_IS_EQUAL(t_result.data()[i], Scalar(0)); + } + } + + // Large contraction with bias addition on zero-sized contraction dimension: + // {64, 0} * {0, 64} -> {64, 64} + { + Tensor<Scalar, 2, DataLayout> t_left_large(64, 0); + Tensor<Scalar, 2, DataLayout> t_right_large(0, 64); + Tensor<Scalar, 2, DataLayout> t_result_large(64, 64); + t_result_large.setConstant(non_zero); + + Eigen::array<DimPair, 1> dims = {{DimPair(1, 0)}}; + t_result_large = t_left_large.contract(t_right_large, dims, AddBiasOutputKernel<Scalar>(bias)); + + VERIFY_IS_EQUAL(t_result_large.dimension(0), 64); + VERIFY_IS_EQUAL(t_result_large.dimension(1), 64); + for (Index i = 0; i < t_result_large.size(); ++i) { + VERIFY_IS_EQUAL(t_result_large.data()[i], bias); + } + } + + // Multidimensional contraction with bias: {5, 0, 4} * {0, 3, 4} -> {5, 3} + { + Tensor<Scalar, 3, DataLayout> t_left(5, 0, 4); + Tensor<Scalar, 3, DataLayout> t_right(0, 3, 4); + Tensor<Scalar, 2, DataLayout> t_result(5, 3); + t_result.setConstant(non_zero); + + Eigen::array<DimPair, 2> dims = {{DimPair(1, 0), DimPair(2, 2)}}; + t_result = t_left.contract(t_right, dims, AddBiasOutputKernel<Scalar>(bias)); + + VERIFY_IS_EQUAL(t_result.dimension(0), 5); + VERIFY_IS_EQUAL(t_result.dimension(1), 3); + for (Index i = 0; i < t_result.dimension(0); ++i) { + for (Index j = 0; j < t_result.dimension(1); ++j) { + VERIFY_IS_EQUAL(t_result(i, j), bias); + } + } + } + + // GEMV-shape contraction with bias (e.g. 1 column in RHS): {10, 0} * {0, 1} -> {10, 1} + { + Tensor<Scalar, 2, DataLayout> t_left(10, 0); + Tensor<Scalar, 2, DataLayout> t_right(0, 1); + Tensor<Scalar, 2, DataLayout> t_result(10, 1); + t_result.setConstant(non_zero); + + Eigen::array<DimPair, 1> dims = {{DimPair(1, 0)}}; + t_result = t_left.contract(t_right, dims, AddBiasOutputKernel<Scalar>(bias)); + + VERIFY_IS_EQUAL(t_result.dimension(0), 10); + VERIFY_IS_EQUAL(t_result.dimension(1), 1); + for (Index i = 0; i < t_result.size(); ++i) { + VERIFY_IS_EQUAL(t_result.data()[i], bias); + } + } +} + EIGEN_DECLARE_TEST(tensor_contraction) { CALL_SUBTEST_1(test_evals<ColMajor>()); CALL_SUBTEST_1(test_evals<RowMajor>()); @@ -637,6 +965,39 @@ CALL_SUBTEST_9((test_narrow_index_contraction<ColMajor, std::complex<double>>())); CALL_SUBTEST_9((test_narrow_index_contraction<RowMajor, std::complex<double>>())); + // Zero-dimension contraction tests: + CALL_SUBTEST_10((test_zero_dim_contraction_2d<ColMajor, float>())); + CALL_SUBTEST_10((test_zero_dim_contraction_2d<RowMajor, float>())); + CALL_SUBTEST_10((test_zero_dim_contraction_2d<ColMajor, double>())); + CALL_SUBTEST_10((test_zero_dim_contraction_2d<RowMajor, double>())); + CALL_SUBTEST_10((test_zero_dim_contraction_2d<ColMajor, std::complex<float>>())); + CALL_SUBTEST_10((test_zero_dim_contraction_2d<RowMajor, std::complex<float>>())); + CALL_SUBTEST_10((test_zero_dim_contraction_2d<ColMajor, std::complex<double>>())); + CALL_SUBTEST_10((test_zero_dim_contraction_2d<RowMajor, std::complex<double>>())); + + CALL_SUBTEST_10((test_zero_dim_contraction_multidims<ColMajor, float>())); + CALL_SUBTEST_10((test_zero_dim_contraction_multidims<RowMajor, float>())); + CALL_SUBTEST_10((test_zero_dim_contraction_multidims<ColMajor, double>())); + CALL_SUBTEST_10((test_zero_dim_contraction_multidims<RowMajor, double>())); + CALL_SUBTEST_10((test_zero_dim_contraction_multidims<ColMajor, std::complex<float>>())); + CALL_SUBTEST_10((test_zero_dim_contraction_multidims<RowMajor, std::complex<float>>())); + CALL_SUBTEST_10((test_zero_dim_contraction_multidims<ColMajor, std::complex<double>>())); + CALL_SUBTEST_10((test_zero_dim_contraction_multidims<RowMajor, std::complex<double>>())); + + CALL_SUBTEST_10((test_zero_dim_contraction_outer_zeros<ColMajor, float>())); + CALL_SUBTEST_10((test_zero_dim_contraction_outer_zeros<RowMajor, float>())); + CALL_SUBTEST_10((test_zero_dim_contraction_outer_zeros<ColMajor, double>())); + CALL_SUBTEST_10((test_zero_dim_contraction_outer_zeros<RowMajor, double>())); + CALL_SUBTEST_10((test_zero_dim_contraction_outer_zeros<ColMajor, std::complex<float>>())); + CALL_SUBTEST_10((test_zero_dim_contraction_outer_zeros<RowMajor, std::complex<float>>())); + CALL_SUBTEST_10((test_zero_dim_contraction_outer_zeros<ColMajor, std::complex<double>>())); + CALL_SUBTEST_10((test_zero_dim_contraction_outer_zeros<RowMajor, std::complex<double>>())); + + CALL_SUBTEST_10((test_zero_dim_contraction_output_kernel<ColMajor, float>())); + CALL_SUBTEST_10((test_zero_dim_contraction_output_kernel<RowMajor, float>())); + CALL_SUBTEST_10((test_zero_dim_contraction_output_kernel<ColMajor, double>())); + CALL_SUBTEST_10((test_zero_dim_contraction_output_kernel<RowMajor, double>())); + // Force CMake to split this test. - // EIGEN_SUFFIXES;1;2;3;4;5;6;7;8;9 + // EIGEN_SUFFIXES;1;2;3;4;5;6;7;8;9;10 }
diff --git a/unsupported/test/tensor_thread_pool.cpp b/unsupported/test/tensor_thread_pool.cpp index c9c7c1e..3f6863d 100644 --- a/unsupported/test/tensor_thread_pool.cpp +++ b/unsupported/test/tensor_thread_pool.cpp
@@ -875,6 +875,263 @@ VERIFY_IS_EQUAL(allocator->dealloc_count(), num_allocs); } +template <int DataLayout> +void test_multithread_zero_dim_contraction() { + const int num_threads = internal::random<int>(2, 11); + ThreadPool threads(num_threads); + Eigen::ThreadPoolDevice device(&threads, num_threads); + + // 2D contraction: {10, 0} * {0, 20} -> {10, 20} + { + Tensor<float, 2, DataLayout> t_left(10, 0); + Tensor<float, 2, DataLayout> t_right(0, 20); + Tensor<float, 2, DataLayout> t_result(10, 20); + t_result.setConstant(123.0f); + + typedef Tensor<float, 1>::DimensionPair DimPair; + Eigen::array<DimPair, 1> dims{{DimPair(1, 0)}}; + + t_result.device(device) = t_left.contract(t_right, dims); + + VERIFY_IS_EQUAL(t_result.dimension(0), 10); + VERIFY_IS_EQUAL(t_result.dimension(1), 20); + for (Index i = 0; i < t_result.size(); ++i) { + VERIFY_IS_EQUAL(t_result.data()[i], 0.0f); + } + } + + // 2D contraction with non-empty dim >= 48: {64, 0} * {0, 64} -> {64, 64} + { + Tensor<float, 2, DataLayout> t_left(64, 0); + Tensor<float, 2, DataLayout> t_right(0, 64); + Tensor<float, 2, DataLayout> t_result(64, 64); + t_result.setConstant(123.0f); + + typedef Tensor<float, 1>::DimensionPair DimPair; + Eigen::array<DimPair, 1> dims{{DimPair(1, 0)}}; + + t_result.device(device) = t_left.contract(t_right, dims); + + VERIFY_IS_EQUAL(t_result.dimension(0), 64); + VERIFY_IS_EQUAL(t_result.dimension(1), 64); + for (Index i = 0; i < t_result.size(); ++i) { + VERIFY_IS_EQUAL(t_result.data()[i], 0.0f); + } + } + + // Multidim contraction: {5, 0, 4} * {0, 3, 4} with dims {1, 0} and {2, 2} -> {5, 3} + { + Tensor<float, 3, DataLayout> t_left(5, 0, 4); + Tensor<float, 3, DataLayout> t_right(0, 3, 4); + Tensor<float, 2, DataLayout> t_result(5, 3); + t_result.setConstant(123.0f); + + typedef Tensor<float, 1>::DimensionPair DimPair; + Eigen::array<DimPair, 2> dims{{DimPair(1, 0), DimPair(2, 2)}}; + + t_result.device(device) = t_left.contract(t_right, dims); + + VERIFY_IS_EQUAL(t_result.dimension(0), 5); + VERIFY_IS_EQUAL(t_result.dimension(1), 3); + for (Index i = 0; i < t_result.size(); ++i) { + VERIFY_IS_EQUAL(t_result.data()[i], 0.0f); + } + } +} + +template <int DataLayout> +void test_async_multithread_zero_dim_contraction() { + const int num_threads = internal::random<int>(2, 11); + ThreadPool threads(num_threads); + Eigen::ThreadPoolDevice device(&threads, num_threads); + + Tensor<float, 2, DataLayout> t_left(10, 0); + Tensor<float, 2, DataLayout> t_right(0, 20); + Tensor<float, 2, DataLayout> t_result(10, 20); + t_result.setConstant(123.0f); + + typedef Tensor<float, 1>::DimensionPair DimPair; + Eigen::array<DimPair, 1> dims{{DimPair(1, 0)}}; + + Eigen::Barrier barrier(1); + t_result.device(device, [&barrier]() { barrier.Notify(); }) = t_left.contract(t_right, dims); + barrier.Wait(); + + VERIFY_IS_EQUAL(t_result.dimension(0), 10); + VERIFY_IS_EQUAL(t_result.dimension(1), 20); + for (Index i = 0; i < t_result.size(); ++i) { + VERIFY_IS_EQUAL(t_result.data()[i], 0.0f); + } +} + +template <typename Scalar> +struct AddBiasOutputKernel { + Scalar bias; + explicit AddBiasOutputKernel(Scalar b) : bias(b) {} + + template <typename Index, typename ResScalar> + EIGEN_ALWAYS_INLINE void operator()(const internal::blas_data_mapper<ResScalar, Index, ColMajor>& output_mapper, + const TensorContractionParams&, Index, Index, Index num_rows, + Index num_cols) const { + for (Index i = 0; i < num_rows; ++i) { + for (Index j = 0; j < num_cols; ++j) { + output_mapper(i, j) += bias; + } + } + } +}; + +template <int DataLayout> +void test_multithread_zero_dim_contraction_with_output_kernel() { + const int num_threads = internal::random<int>(2, 11); + ThreadPool threads(num_threads); + Eigen::ThreadPoolDevice device(&threads, num_threads); + + const float bias = 42.0f; + typedef Tensor<float, 1>::DimensionPair DimPair; + Eigen::array<DimPair, 1> dims{{DimPair(1, 0)}}; + + // {10, 0} * {0, 20} -> {10, 20} + { + Tensor<float, 2, DataLayout> t_left(10, 0); + Tensor<float, 2, DataLayout> t_right(0, 20); + Tensor<float, 2, DataLayout> t_result(10, 20); + t_result.setConstant(123.0f); + + t_result.device(device) = t_left.contract(t_right, dims, AddBiasOutputKernel<float>(bias)); + + VERIFY_IS_EQUAL(t_result.dimension(0), 10); + VERIFY_IS_EQUAL(t_result.dimension(1), 20); + for (Index i = 0; i < t_result.size(); ++i) { + VERIFY_IS_EQUAL(t_result.data()[i], bias); + } + } + + // {64, 0} * {0, 64} -> {64, 64} + { + Tensor<float, 2, DataLayout> t_left(64, 0); + Tensor<float, 2, DataLayout> t_right(0, 64); + Tensor<float, 2, DataLayout> t_result(64, 64); + t_result.setConstant(123.0f); + + t_result.device(device) = t_left.contract(t_right, dims, AddBiasOutputKernel<float>(bias)); + + VERIFY_IS_EQUAL(t_result.dimension(0), 64); + VERIFY_IS_EQUAL(t_result.dimension(1), 64); + for (Index i = 0; i < t_result.size(); ++i) { + VERIFY_IS_EQUAL(t_result.data()[i], bias); + } + } +} + +template <int DataLayout> +void test_async_multithread_zero_dim_contraction_with_output_kernel() { + const int num_threads = internal::random<int>(2, 11); + ThreadPool threads(num_threads); + Eigen::ThreadPoolDevice device(&threads, num_threads); + + const float bias = 42.0f; + typedef Tensor<float, 1>::DimensionPair DimPair; + Eigen::array<DimPair, 1> dims{{DimPair(1, 0)}}; + + // {10, 0} * {0, 20} -> {10, 20} + { + Tensor<float, 2, DataLayout> t_left(10, 0); + Tensor<float, 2, DataLayout> t_right(0, 20); + Tensor<float, 2, DataLayout> t_result(10, 20); + t_result.setConstant(123.0f); + + Eigen::Barrier barrier(1); + t_result.device(device, [&barrier]() { barrier.Notify(); }) = + t_left.contract(t_right, dims, AddBiasOutputKernel<float>(bias)); + barrier.Wait(); + + VERIFY_IS_EQUAL(t_result.dimension(0), 10); + VERIFY_IS_EQUAL(t_result.dimension(1), 20); + for (Index i = 0; i < t_result.size(); ++i) { + VERIFY_IS_EQUAL(t_result.data()[i], bias); + } + } + + // {64, 0} * {0, 64} -> {64, 64} + { + Tensor<float, 2, DataLayout> t_left(64, 0); + Tensor<float, 2, DataLayout> t_right(0, 64); + Tensor<float, 2, DataLayout> t_result(64, 64); + t_result.setConstant(123.0f); + + Eigen::Barrier barrier(1); + t_result.device(device, [&barrier]() { barrier.Notify(); }) = + t_left.contract(t_right, dims, AddBiasOutputKernel<float>(bias)); + barrier.Wait(); + + VERIFY_IS_EQUAL(t_result.dimension(0), 64); + VERIFY_IS_EQUAL(t_result.dimension(1), 64); + for (Index i = 0; i < t_result.size(); ++i) { + VERIFY_IS_EQUAL(t_result.data()[i], bias); + } + } +} + +template <int DataLayout> +void test_multithread_zero_dim_outer_zeros() { + const int num_threads = internal::random<int>(2, 11); + ThreadPool threads(num_threads); + Eigen::ThreadPoolDevice device(&threads, num_threads); + + typedef Tensor<float, 1>::DimensionPair DimPair; + Eigen::array<DimPair, 1> dims{{DimPair(1, 0)}}; + + // Outer zero on LHS: {0, 5} * {5, 10} -> {0, 10} + { + Tensor<float, 2, DataLayout> t_left(0, 5); + Tensor<float, 2, DataLayout> t_right(5, 10); + t_right.setRandom(); + Tensor<float, 2, DataLayout> t_result(0, 10); + t_result.device(device) = t_left.contract(t_right, dims); + VERIFY_IS_EQUAL(t_result.dimension(0), 0); + VERIFY_IS_EQUAL(t_result.dimension(1), 10); + VERIFY_IS_EQUAL(t_result.size(), 0); + } + + // Outer zero on RHS: {10, 5} * {5, 0} -> {10, 0} + { + Tensor<float, 2, DataLayout> t_left(10, 5); + Tensor<float, 2, DataLayout> t_right(5, 0); + t_left.setRandom(); + Tensor<float, 2, DataLayout> t_result(10, 0); + t_result.device(device) = t_left.contract(t_right, dims); + VERIFY_IS_EQUAL(t_result.dimension(0), 10); + VERIFY_IS_EQUAL(t_result.dimension(1), 0); + VERIFY_IS_EQUAL(t_result.size(), 0); + } + + // Outer zero on both: {0, 5} * {5, 0} -> {0, 0} + { + Tensor<float, 2, DataLayout> t_left(0, 5); + Tensor<float, 2, DataLayout> t_right(5, 0); + Tensor<float, 2, DataLayout> t_result(0, 0); + t_result.device(device) = t_left.contract(t_right, dims); + VERIFY_IS_EQUAL(t_result.dimension(0), 0); + VERIFY_IS_EQUAL(t_result.dimension(1), 0); + VERIFY_IS_EQUAL(t_result.size(), 0); + } + + // Async zero outer dim + { + Tensor<float, 2, DataLayout> t_left(0, 5); + Tensor<float, 2, DataLayout> t_right(5, 10); + t_right.setRandom(); + Tensor<float, 2, DataLayout> t_result(0, 10); + Eigen::Barrier barrier(1); + t_result.device(device, [&barrier]() { barrier.Notify(); }) = t_left.contract(t_right, dims); + barrier.Wait(); + VERIFY_IS_EQUAL(t_result.dimension(0), 0); + VERIFY_IS_EQUAL(t_result.dimension(1), 10); + VERIFY_IS_EQUAL(t_result.size(), 0); + } +} + EIGEN_DECLARE_TEST(tensor_thread_pool) { CALL_SUBTEST_1(test_multithread_elementwise()); CALL_SUBTEST_1(test_async_multithread_elementwise()); @@ -932,6 +1189,17 @@ CALL_SUBTEST_13(test_multithread_shuffle<RowMajor>(&test_allocator)); CALL_SUBTEST_13(test_threadpool_allocate(&test_allocator)); + CALL_SUBTEST_14(test_multithread_zero_dim_contraction<ColMajor>()); + CALL_SUBTEST_14(test_multithread_zero_dim_contraction<RowMajor>()); + CALL_SUBTEST_14(test_async_multithread_zero_dim_contraction<ColMajor>()); + CALL_SUBTEST_14(test_async_multithread_zero_dim_contraction<RowMajor>()); + CALL_SUBTEST_14(test_multithread_zero_dim_contraction_with_output_kernel<ColMajor>()); + CALL_SUBTEST_14(test_multithread_zero_dim_contraction_with_output_kernel<RowMajor>()); + CALL_SUBTEST_14(test_async_multithread_zero_dim_contraction_with_output_kernel<ColMajor>()); + CALL_SUBTEST_14(test_async_multithread_zero_dim_contraction_with_output_kernel<RowMajor>()); + CALL_SUBTEST_14(test_multithread_zero_dim_outer_zeros<ColMajor>()); + CALL_SUBTEST_14(test_multithread_zero_dim_outer_zeros<RowMajor>()); + // Force CMake to split this test. - // EIGEN_SUFFIXES;1;2;3;4;5;6;7;8;9;10;11;12;13 + // EIGEN_SUFFIXES;1;2;3;4;5;6;7;8;9;10;11;12;13;14 }