Fix Tensor thread-pool cost model, shuffling tiling, and contraction bandwidth libeigen/eigen!3367 Co-authored-by: Antonio Sanchez <cantonios@google.com>
diff --git a/contrib/Eigen/src/Tensor/TensorContractionThreadPool.h b/contrib/Eigen/src/Tensor/TensorContractionThreadPool.h index 96da0ca..733c020 100644 --- a/contrib/Eigen/src/Tensor/TensorContractionThreadPool.h +++ b/contrib/Eigen/src/Tensor/TensorContractionThreadPool.h
@@ -1450,7 +1450,7 @@ const int packed_size = std::min<int>(PacketType<LhsScalar, Device>::size, PacketType<RhsScalar, Device>::size); const int output_packet_size = internal::unpacket_traits<PacketReturnType>::size; const double kd = static_cast<double>(bk); - double compute_bandwidth = computeBandwidth(false, bm, bn, bk); + double compute_bandwidth = computeBandwidth(bm, bn, bk); // Computations. TensorOpCost cost = TensorOpCost(0, 0, kd * compute_bandwidth, true, packed_size); // Output stores. @@ -1499,7 +1499,7 @@ TensorOpCost contractionCostPerInnerDim(Index m, Index n, Index k) const { // Compute cost. const int output_packet_size = internal::unpacket_traits<PacketReturnType>::size; - TensorOpCost cost(0, 0, (computeBandwidth(true, m, n, k) * m) * n, true, output_packet_size); + TensorOpCost cost(0, 0, (computeBandwidth(m, n, k) * m) * n, true, output_packet_size); // Output stores. cost += TensorOpCost(0, sizeof(CoeffReturnType), 0, true, output_packet_size); TensorOpCost lhsCost = this->m_leftImpl.costPerCoeff(true) * m; @@ -1533,13 +1533,11 @@ return num_threads; } - double computeBandwidth(bool shard_by_col, Index bm, Index bn, Index bk) const { + double computeBandwidth(Index bm, Index bn, Index bk) const { // Peak VFMA bandwidth is 0.5. However if we have not enough data for // vectorization bandwidth drops. The 4.0 and 2.0 bandwidth is determined // experimentally. - double computeBandwidth = bk == 1 ? 4.0 - : (shard_by_col ? bn : bm) < Traits::nr || (shard_by_col ? bm : bn) < Traits::mr ? 2.0 - : 0.5; + double computeBandwidth = bk == 1 ? 4.0 : (bn < Traits::nr || bm < Traits::mr) ? 2.0 : 0.5; #ifndef EIGEN_VECTORIZE_FMA // Bandwidth of all of VFMA/MULPS/ADDPS is 0.5 on latest Intel processors. // However for MULPS/ADDPS we have dependent sequence of 2 such
diff --git a/contrib/Eigen/src/Tensor/TensorCostModel.h b/contrib/Eigen/src/Tensor/TensorCostModel.h index 8b8a7ea..c2e1626 100644 --- a/contrib/Eigen/src/Tensor/TensorCostModel.h +++ b/contrib/Eigen/src/Tensor/TensorCostModel.h
@@ -145,11 +145,11 @@ // Scaling from Eigen compute cost to device cycles. static constexpr int kDeviceCyclesPerComputeCycle = 1; - // Thread overhead in device cycles. ~8us at 3GHz. + // Thread overhead in device cycles. ~33us at 3GHz. // Minimum total work to justify thread pool dispatch. - static constexpr int kStartupCycles = 25000; + static constexpr int kStartupCycles = 100000; // Minimum work per thread to amortize dispatch and synchronization overhead. - static constexpr int kPerThreadCycles = 25000; + static constexpr int kPerThreadCycles = 50000; static constexpr int kTaskSize = 40000; // Memory bandwidth saturation: on typical multi-socket servers, 2-6 cores
diff --git a/contrib/Eigen/src/Tensor/TensorEvaluator.h b/contrib/Eigen/src/Tensor/TensorEvaluator.h index a04a528..84decf9 100644 --- a/contrib/Eigen/src/Tensor/TensorEvaluator.h +++ b/contrib/Eigen/src/Tensor/TensorEvaluator.h
@@ -641,7 +641,8 @@ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE internal::TensorBlockResourceRequirements getResourceRequirements() const { static constexpr double functor_cost = internal::functor_traits<UnaryOp>::Cost; - return m_argImpl.getResourceRequirements().addCostPerCoeff({0, 0, functor_cost / PacketSize}); + return m_argImpl.getResourceRequirements().addCostPerCoeff( + TensorOpCost(0, 0, functor_cost, PacketAccess, PacketSize)); } EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlock block(TensorBlockDesc& desc, TensorBlockScratch& scratch, @@ -768,7 +769,7 @@ static constexpr double functor_cost = internal::functor_traits<BinaryOp>::Cost; return internal::TensorBlockResourceRequirements::merge(m_leftImpl.getResourceRequirements(), m_rightImpl.getResourceRequirements()) - .addCostPerCoeff({0, 0, functor_cost / PacketSize}); + .addCostPerCoeff(TensorOpCost(0, 0, functor_cost, PacketAccess, PacketSize)); } EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlock block(TensorBlockDesc& desc, TensorBlockScratch& scratch, @@ -912,7 +913,7 @@ m_arg1Impl.getResourceRequirements(), internal::TensorBlockResourceRequirements::merge(m_arg2Impl.getResourceRequirements(), m_arg3Impl.getResourceRequirements())) - .addCostPerCoeff({0, 0, functor_cost / PacketSize}); + .addCostPerCoeff(TensorOpCost(0, 0, functor_cost, PacketAccess, PacketSize)); } EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlock block(TensorBlockDesc& desc, TensorBlockScratch& scratch,
diff --git a/contrib/Eigen/src/Tensor/TensorExecutor.h b/contrib/Eigen/src/Tensor/TensorExecutor.h index eacd35f..994abc5 100644 --- a/contrib/Eigen/src/Tensor/TensorExecutor.h +++ b/contrib/Eigen/src/Tensor/TensorExecutor.h
@@ -186,9 +186,11 @@ // Query expression tree for desired block size/shape. TensorBlockResourceRequirements requirements = evaluator.getResourceRequirements(); - // Update target block size based on cost model. - double taskSize = TensorCostModel<ThreadPoolDevice>::taskSize(1, requirements.cost_per_coeff); - requirements.size = static_cast<size_t>(1.0 / taskSize); + // Update target block size based on cost model, preserving tighter + // cache-targeted block size bounds requested by subexpressions. + double task_inverse_ratio = TensorCostModel<ThreadPoolDevice>::taskSize(1, requirements.cost_per_coeff); + const size_t task_size = numext::maxi<size_t>(1, static_cast<size_t>(1.0 / task_inverse_ratio)); + requirements.size = (requirements.size > 1) ? numext::mini(requirements.size, task_size) : task_size; TensorBlockMapper block_mapper(typename TensorBlockMapper::Dimensions(evaluator.dimensions()), requirements);
diff --git a/contrib/Eigen/src/Tensor/TensorShuffling.h b/contrib/Eigen/src/Tensor/TensorShuffling.h index 8c203e7..fb901b8 100644 --- a/contrib/Eigen/src/Tensor/TensorShuffling.h +++ b/contrib/Eigen/src/Tensor/TensorShuffling.h
@@ -256,14 +256,14 @@ const size_t target_size = m_device.firstLevelCacheSize(); const bool inner_dim_shuffled = m_shuffle[inner_dim] != inner_dim; - // Shuffled inner dimensions lead to a random memory access, which is not - // captured by default cost model bytes loaded/stored. We add this cost - // explicitly. The number of cycles was picked based on the benchmarks. - // TODO(ezhulenev): This number was picked based on very questionable - // benchmarks, add benchmarks that are representative of real workloads. + // Shuffling the inner dimension requires a strided gather across source + // rows into the block buffer, which is unvectorized across lanes (~1 cycle + // per scalar element in pgather) and spans one cache line per source row + // in the tile. Use target_size / 8 (~4 KiB) so the gathered source cache + // lines stay within L1 cache capacity and set associativity. using BlockRequirements = internal::TensorBlockResourceRequirements; if (inner_dim_shuffled) { - return BlockRequirements::uniform<Scalar>(target_size).addCostPerCoeff({0, 0, NumDims * 28}); + return BlockRequirements::uniform<Scalar>(target_size / 8).addCostPerCoeff({0, 0, 1.0}); } else { return BlockRequirements::skewed<Scalar>(target_size); } @@ -296,8 +296,7 @@ ? TensorOpCost::AddCost<Index>() : NumDims * (2 * TensorOpCost::AddCost<Index>() + 2 * TensorOpCost::MulCost<Index>() + TensorOpCost::DivCost<Index>()); - return m_impl.costPerCoeff(vectorized) + - TensorOpCost(0, 0, compute_cost, m_is_identity /* vectorized */, PacketSize); + return m_impl.costPerCoeff(vectorized) + TensorOpCost(0, 0, compute_cost, vectorized, PacketSize); } EIGEN_DEVICE_FUNC typename Storage::Type data() const { return nullptr; }
diff --git a/contrib/benchmarks/Tensor/bench_contraction.cpp b/contrib/benchmarks/Tensor/bench_contraction.cpp index cd36154..88e9e97 100644 --- a/contrib/benchmarks/Tensor/bench_contraction.cpp +++ b/contrib/benchmarks/Tensor/bench_contraction.cpp
@@ -146,5 +146,5 @@ BENCHMARK(BM_Contraction) CONTRACTION_SIZES; BENCHMARK(BM_Contraction_RowMajor) CONTRACTION_SIZES; -BENCHMARK(BM_Contraction_ThreadPool) CONTRACTION_THREADPOOL_SIZES; +BENCHMARK(BM_Contraction_ThreadPool) CONTRACTION_THREADPOOL_SIZES->UseRealTime(); BENCHMARK(BM_BatchContraction) BATCH_SIZES;
diff --git a/contrib/test/tensor_cost_model.cpp b/contrib/test/tensor_cost_model.cpp index e54f822..c6fc9e5 100644 --- a/contrib/test/tensor_cost_model.cpp +++ b/contrib/test/tensor_cost_model.cpp
@@ -57,17 +57,121 @@ static void test_cost_model_constants_odr_use() { using Cost = Eigen::TensorCostModel<DefaultDevice>; VERIFY_IS_EQUAL(read_through_address(Cost::kDeviceCyclesPerComputeCycle), 1); - VERIFY_IS_EQUAL(read_through_address(Cost::kStartupCycles), 25000); - VERIFY_IS_EQUAL(read_through_address(Cost::kPerThreadCycles), 25000); + VERIFY_IS_EQUAL(read_through_address(Cost::kStartupCycles), 100000); + VERIFY_IS_EQUAL(read_through_address(Cost::kPerThreadCycles), 50000); VERIFY_IS_EQUAL(read_through_address(Cost::kTaskSize), 40000); VERIFY_IS_EQUAL(read_through_address(Cost::kMemBandwidthSaturationThreads), 4); VERIFY_IS_EQUAL(read_through_address(Cost::kMemBoundThreshold), 2.0); VERIFY_IS_EQUAL(read_through_address(Cost::kDramThresholdBytes), 1024.0 * 1024.0); } +template <typename Scalar> +struct scalar_only_unary_op { + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar operator()(const Scalar& x) const { return -x; } +}; + +template <typename Scalar> +struct scalar_only_binary_op { + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar operator()(const Scalar& x, const Scalar& y) const { return x + y; } +}; + +template <typename Scalar> +struct scalar_only_ternary_op { + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar operator()(const Scalar& x, const Scalar& y, const Scalar& z) const { + return x + y + z; + } +}; + +namespace Eigen { +namespace internal { +template <typename Scalar> +struct functor_traits<scalar_only_unary_op<Scalar>> { + enum { Cost = 16, PacketAccess = false }; +}; +template <typename Scalar> +struct functor_traits<scalar_only_binary_op<Scalar>> { + enum { Cost = 16, PacketAccess = false }; +}; +template <typename Scalar> +struct functor_traits<scalar_only_ternary_op<Scalar>> { + enum { Cost = 16, PacketAccess = false }; +}; +} // namespace internal +} // namespace Eigen + +template <typename Scalar> +static void test_unary_binary_ternary_block_resource_cost() { + Tensor<Scalar, 2> a(64, 64); + Tensor<Scalar, 2> b(64, 64); + Tensor<Scalar, 2> c(64, 64); + a.setRandom(); + b.setRandom(); + c.setRandom(); + + DefaultDevice device; + + auto unary_expr = a.unaryExpr(scalar_only_unary_op<Scalar>()); + using UnaryEval = TensorEvaluator<const decltype(unary_expr), DefaultDevice>; + STATIC_CHECK(!UnaryEval::PacketAccess); + UnaryEval unary_eval(unary_expr, device); + VERIFY_IS_EQUAL(unary_eval.getResourceRequirements().cost_per_coeff.compute_cycles(), 16.0); + + auto binary_expr = a.binaryExpr(b, scalar_only_binary_op<Scalar>()); + using BinaryEval = TensorEvaluator<const decltype(binary_expr), DefaultDevice>; + STATIC_CHECK(!BinaryEval::PacketAccess); + BinaryEval binary_eval(binary_expr, device); + VERIFY_IS_EQUAL(binary_eval.getResourceRequirements().cost_per_coeff.compute_cycles(), 16.0); + + Eigen::TensorCwiseTernaryOp<scalar_only_ternary_op<Scalar>, const decltype(a), const decltype(b), const decltype(c)> + ternary_expr(a, b, c, scalar_only_ternary_op<Scalar>()); + using TernaryEval = TensorEvaluator<const decltype(ternary_expr), DefaultDevice>; + STATIC_CHECK(!TernaryEval::PacketAccess); + TernaryEval ternary_eval(ternary_expr, device); + VERIFY_IS_EQUAL(ternary_eval.getResourceRequirements().cost_per_coeff.compute_cycles(), 16.0); +} + +template <typename Scalar> +static void test_shuffling_cost_and_block_requirements() { + Tensor<Scalar, 2> a(256, 256); + a.setRandom(); + + Eigen::array<Eigen::Index, 2> perm = {1, 0}; + auto shuffle_expr = a.shuffle(perm); + using ShuffleEval = TensorEvaluator<const decltype(shuffle_expr), DefaultDevice>; + + DefaultDevice device; + ShuffleEval eval(shuffle_expr, device); + + const auto scalar_cost = eval.costPerCoeff(/*vectorized=*/false); + const auto vec_cost = eval.costPerCoeff(/*vectorized=*/true); + if (ShuffleEval::PacketAccess && ShuffleEval::PacketSize > 1) { + VERIFY(vec_cost.compute_cycles() < scalar_cost.compute_cycles()); + } + + const auto req = eval.getResourceRequirements(); + VERIFY(req.shape_type == Eigen::internal::TensorBlockShapeType::kUniformAllDims); + VERIFY_IS_EQUAL(req.cost_per_coeff.bytes_loaded(), static_cast<double>(sizeof(Scalar))); + VERIFY_IS_EQUAL(req.cost_per_coeff.bytes_stored(), static_cast<double>(sizeof(Scalar))); + VERIFY_IS_EQUAL(req.cost_per_coeff.compute_cycles(), 1.0); + // Small 256x256 2D transpose should stay single-threaded rather than paying + // thread-pool dispatch and barrier overhead. + const int threads_256 = + Eigen::TensorCostModel<DefaultDevice>::numThreads(256.0 * 256.0, req.cost_per_coeff, /*max_threads=*/8); + VERIFY_IS_EQUAL(threads_256, 1); + + // Larger 1024x1024 2D transpose should parallelize across threads. + const int threads_1024 = + Eigen::TensorCostModel<DefaultDevice>::numThreads(1024.0 * 1024.0, req.cost_per_coeff, /*max_threads=*/8); + VERIFY(threads_1024 > 1); +} + EIGEN_DECLARE_TEST(tensor_cost_model) { CALL_SUBTEST(test_cost_model_constants_odr_use()); CALL_SUBTEST(test_nullary_zero_bytes_loaded<float>()); CALL_SUBTEST(test_nullary_zero_bytes_loaded<double>()); CALL_SUBTEST(test_nullary_zero_bytes_loaded<int>()); + CALL_SUBTEST(test_unary_binary_ternary_block_resource_cost<float>()); + CALL_SUBTEST(test_unary_binary_ternary_block_resource_cost<double>()); + CALL_SUBTEST(test_shuffling_cost_and_block_requirements<float>()); + CALL_SUBTEST(test_shuffling_cost_and_block_requirements<double>()); }