Repository navigation
Tensor: host ops cover all batches of a batched tensor - #614
Merged
Merged
Conversation
Value-returning ops (clone, scale, neg, add, subt, mult, unary, binary, with or without permutation) returned a single batch because the op-constructors hard-coded nbatch = 1; in-place ops (scale_to, neg_to, add_to, ..., inplace_unary) touched only batch 0 because the contiguous inplace_tensor_op/tensor_init kernels swept range().volume() elements; permutation threw on at_ordinal's nbatch() == 1 assertion. - op-constructors take the batch count from their (left) argument - contiguous kernels sweep volume * nbatch and assert batch congruence - detail::permute permutes each batch separately - add detail::nbatch, is_batched_tensor_v, is_nbatch_congruent helpers - test tensor_suite/batched_ops
…ll batches - check_nbatch_congruent (TA_EXCEPTION, all builds) in every inplace_tensor_op/tensor_init kernel and in detail::permute; the non-contiguous kernels previously had no check, so a batched tensor combined with a view returned uninitialized batches, and in Release a mismatch read/wrote past the smaller operand - Tensor(other, op, perm) and Tensor(left, right, op, perm) apply the inner permutation of a tensor of tensors to all batches, not just the first (range-for stops at size()) - document which Tensor accessors describe only the first batch - tests: each op checked independently; mismatched and view operands are rejected; batched tensor-of-tensors ops incl. inner permutation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Host
TA::Tensorops ignored all but the first batch of a batched tensor (nbatch() > 1):clone,scale,neg,add,subt,mult,unary,binary(with or without a permutation) go through the op-constructorsTensor(other, op[, perm])/Tensor(left, right, op[, perm]), which hard-codednbatch = 1.scale_to,neg_to,add_to,subt_to,mult_to,inplace_unaryetc. go through the contiguousinplace_tensor_opkernel, which sweptrange().volume()elements.permutekernel addresses elements viaat_ordinal, which assertsnbatch() == 1.Reductions (
sum,squared_norm,dot, ...) already usedtotal_size()and were correct.Changes
tensor/utility.h:detail::nbatch(t)(1 for types withoutnbatch()),detail::is_batched_tensor_v<T>(hasnbatch()andbatch(idx)),detail::is_nbatch_congruent(ts...), anddetail::check_nbatch_congruent(ts...), which throwsTA_EXCEPTIONin all builds. A comment lists whichTensoraccessors describe only the first batch (range(),size(),begin()/end(),operator[],at_ordinal).tensor/kernels.h: the contiguousinplace_tensor_opandtensor_initkernels sweeprange().volume() * nbatchelements. Everyinplace_tensor_op/tensor_initkernel, contiguous or not, rejects operands with different batch counts.tensor/permute.h:detail::permuteloops over batches, permuting eachbatch(b)(the range describes one batch), and rejects mismatched batch counts.tensor/tensor.h: the op-constructors and the converting constructor take the batch count from their (left) argument. For a tensor of tensors,Tensor(other, op, perm)andTensor(left, right, op, perm)apply the inner permutation to every batch.Behavior for
nbatch() == 1tensors is unchanged.Behavior change: combining tensors with different batch counts now throws
TiledArray::Exceptionin every build. Before, such a call silently processed only the first batch. The check is a hard check rather than aTA_ASSERTbecause with whole-tensor sweeps a mismatch would otherwise read or write past the end of the smaller operand in Release. A batched tensor combined with a non-contiguous view (TensorInterface) is likewise rejected; without the check it returned batches that were never written.Testing
New tests in
tensor_suite:batched_opschecks every op above on a 3-batchTensor<double>against the op applied batch-by-batch to plain copies, plus the reductions. Each op is checked independently, so one throwing op doesn't hide the rest. On master it fails (clonereturns 1 batch;permutethrows). It also checks that mismatched batch counts and view operands throw.batched_tot_opsdoes the same for a batchedTensor<Tensor<double>>:clone,scale,permuteandaddwith a bipartite permutation, and the op+permutation constructor. With the old inner-permutation loop it fails.Local macOS Debug build (
TA_ASSERT_THROW, Pthreads): fullta_testpasses at np=1 (!@distributed) and np=2 (!@serial).Found while adding batched-tile tests for the device path in #554, whose device branches already cover all batches.
🤖 Generated with Claude Code
https://claude.ai/code/session_01VUXmaoaBxkXhMkXUT1Vn9v