Prune unreachable rhs_inner_dim_reordered=true contraction instantiations libeigen/eigen!3383 Co-authored-by: zce <zce@google.com>
diff --git a/contrib/Eigen/src/Tensor/TensorContraction.h b/contrib/Eigen/src/Tensor/TensorContraction.h index e3e1f8c..2dfb390 100644 --- a/contrib/Eigen/src/Tensor/TensorContraction.h +++ b/contrib/Eigen/src/Tensor/TensorContraction.h
@@ -586,6 +586,7 @@ m_rhs_inner_dim_contiguous = false; } } + eigen_assert(!m_rhs_inner_dim_reordered || (ContractDims > 1 && !m_rhs_inner_dim_contiguous)); // After the sort above, eval_op_indices[*].first is in ascending order. The // contracted LHS dims are the contiguous prefix of storage iff that prefix @@ -658,6 +659,23 @@ } #endif // EIGEN_USE_THREADS + template <typename Func> + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE static void dispatchContraction(Func&& fn, bool lhs_inner_dim_contiguous, + bool rhs_inner_dim_contiguous, + bool rhs_inner_dim_reordered) { + // m_rhs_inner_dim_reordered is only ever set when ContractDims > 1 and + // !m_rhs_inner_dim_contiguous. Canonicalizing rhs_r at compile time avoids + // instantiating unreachable contractions (reducing 8 combinations to 4 when + // ContractDims <= 1, and to 6 when ContractDims > 1). + internal::tensor_contraction_dispatch( + [&](auto lhs_c, auto rhs_c, auto rhs_r) { + constexpr bool kRhsReordered = (ContractDims > 1) && !rhs_c() && rhs_r(); + fn(lhs_c, rhs_c, internal::bool_constant<kRhsReordered>{}); + }, + lhs_inner_dim_contiguous, rhs_inner_dim_contiguous, + (ContractDims > 1) && !rhs_inner_dim_contiguous && rhs_inner_dim_reordered); + } + 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) { @@ -1058,7 +1076,7 @@ template <int Alignment> void evalProduct(Scalar* buffer) const { - internal::tensor_contraction_dispatch( + Base::dispatchContraction( [&](auto lhs_c, auto rhs_c, auto rhs_r) { this->template evalProductSequential<lhs_c(), rhs_c(), rhs_r(), Alignment>(buffer); },
diff --git a/contrib/Eigen/src/Tensor/TensorContractionThreadPool.h b/contrib/Eigen/src/Tensor/TensorContractionThreadPool.h index 2b40248..b09a252 100644 --- a/contrib/Eigen/src/Tensor/TensorContractionThreadPool.h +++ b/contrib/Eigen/src/Tensor/TensorContractionThreadPool.h
@@ -83,6 +83,8 @@ evalProductImpl<EvalToCallback, Alignment>(buffer, std::move(done)); } + using Base::dispatchContraction; + template <typename DoneCallback, int Alignment> void evalProductImpl(Scalar* buffer, DoneCallback done) const { // This function computes a lot of heuristics in multiple steps, and it @@ -110,7 +112,7 @@ return; } if (k == 0) { - internal::tensor_contraction_dispatch( + dispatchContraction( [&](auto lhs_c, auto rhs_c, auto rhs_r) { this->template evalProductSequential<lhs_c(), rhs_c(), rhs_r(), Unaligned>(buffer); }, @@ -189,7 +191,7 @@ if (n == 1) num_threads = 1; if (num_threads == 1) { - internal::tensor_contraction_dispatch( + dispatchContraction( [&](auto lhs_c, auto rhs_c, auto rhs_r) { this->template evalProductSequential<lhs_c(), rhs_c(), rhs_r(), Unaligned>(buffer); }, @@ -275,15 +277,15 @@ // optimization. if (parallelize_by_sharding_dim_only) parallel_pack = false; - internal::tensor_contraction_dispatch( + dispatchContraction( [&](auto lhs_c, auto rhs_c, auto rhs_r) { EIGEN_IF_CONSTEXPR (IsEvalInSyncMode) { - EvalParallelContext<NoCallback, lhs_c(), rhs_c(), rhs_r(), Alignment> ctx( + EvalParallelContext<NoCallback, lhs_c(), rhs_c(), rhs_r()> ctx( this, num_threads, buffer, m, n, k, bm, bn, bk, nm, nn, nk, gm, gn, nm0, nn0, shard_by_col, parallel_pack, parallelize_by_sharding_dim_only, NoCallback()); ctx.run(); } else { - auto* ctx = new EvalParallelContext<DoneCallback, lhs_c(), rhs_c(), rhs_r(), Alignment>( + auto* ctx = new EvalParallelContext<DoneCallback, lhs_c(), rhs_c(), rhs_r()>( this, num_threads, buffer, m, n, k, bm, bn, bk, nm, nn, nk, gm, gn, nm0, nn0, shard_by_col, parallel_pack, parallelize_by_sharding_dim_only, std::move(done)); ctx->run(); @@ -348,7 +350,7 @@ // accessible by block packing and kernel tasks. template <typename DoneCallback, bool lhs_inner_dim_contiguous, bool rhs_inner_dim_contiguous, - bool rhs_inner_dim_reordered, int Alignment> + bool rhs_inner_dim_reordered> class EvalParallelContext { public: using LhsMapper = @@ -1154,7 +1156,7 @@ void processBlock(Index block_idx, Index begin, Index end) { Scalar* buf = block_buffers[block_idx]; - internal::tensor_contraction_dispatch( + dispatchContraction( [&](auto lhs_c, auto rhs_c, auto rhs_r) { evaluator->template evalGemmPartialWithoutOutputKernel<lhs_c(), rhs_c(), rhs_r(), Alignment>( buf, begin, end, /*num_threads=*/internal::convert_index<int>(num_blocks));
diff --git a/contrib/test/tensor_contraction.cpp b/contrib/test/tensor_contraction.cpp index e6c2d33..b9c3aa4 100644 --- a/contrib/test/tensor_contraction.cpp +++ b/contrib/test/tensor_contraction.cpp
@@ -292,6 +292,44 @@ mat1(0, 1, 1) * mat2(1, 1, 0) + mat1(1, 1, 1) * mat2(1, 1, 1)); } +template <int DataLayout, int ContractDims> +static void check_dispatch_pruning_for_rank() { + using Tensor3 = Tensor<float, 3, DataLayout>; + using Contraction = TensorContractionOp<const Eigen::array<DimPair, ContractDims>, const Tensor3, const Tensor3>; + using Evaluator = TensorEvaluator<const Contraction, DefaultDevice>; + + for (int lhs_c_int = 0; lhs_c_int < 2; ++lhs_c_int) { + for (int rhs_c_int = 0; rhs_c_int < 2; ++rhs_c_int) { + for (int rhs_r_int = 0; rhs_r_int < 2; ++rhs_r_int) { + const bool lhs_c_in = lhs_c_int != 0; + const bool rhs_c_in = rhs_c_int != 0; + const bool rhs_r_in = rhs_r_int != 0; + bool dispatched_lhs_c = false; + bool dispatched_rhs_c = false; + bool dispatched_rhs_r = false; + Evaluator::dispatchContraction( + [&](auto lhs_c, auto rhs_c, auto rhs_r) { + static_assert(!rhs_r() || (ContractDims > 1 && !rhs_c()), "unreachable contraction instantiated"); + dispatched_lhs_c = lhs_c(); + dispatched_rhs_c = rhs_c(); + dispatched_rhs_r = rhs_r(); + }, + lhs_c_in, rhs_c_in, rhs_r_in); + VERIFY_IS_EQUAL(dispatched_lhs_c, lhs_c_in); + VERIFY_IS_EQUAL(dispatched_rhs_c, rhs_c_in); + VERIFY_IS_EQUAL(dispatched_rhs_r, (ContractDims > 1) && !rhs_c_in && rhs_r_in); + } + } + } +} + +template <int DataLayout> +static void test_dispatch_pruning() { + check_dispatch_pruning_for_rank<DataLayout, 0>(); + check_dispatch_pruning_for_rank<DataLayout, 1>(); + check_dispatch_pruning_for_rank<DataLayout, 2>(); +} + template <int DataLayout> static void test_consistency() { // this does something like testing (A*B)^T = (B^T * A^T) @@ -955,6 +993,8 @@ CALL_SUBTEST_4(test_expr<RowMajor>()); CALL_SUBTEST_4(test_out_of_order_contraction<ColMajor>()); CALL_SUBTEST_4(test_out_of_order_contraction<RowMajor>()); + CALL_SUBTEST_4(test_dispatch_pruning<ColMajor>()); + CALL_SUBTEST_4(test_dispatch_pruning<RowMajor>()); CALL_SUBTEST_5(test_consistency<ColMajor>()); CALL_SUBTEST_5(test_consistency<RowMajor>()); CALL_SUBTEST_5(test_large_contraction<ColMajor>());