[ajay-mk]

Summary

Adds device (CUDA/HIP) support to TA::Tensor through TA::UMTensor<T>, a TA::Tensor backed by Umpire-managed unified memory (Tensor<T, device_um_allocator<T>>).

Following the review discussion, the device code now lives inside TA::Tensor itself instead of a separate set of UMTensor overloads: each tile op checks whether its operands are UM tensors of a BLAS element type (float, double and their complex types) and, if so, runs a device kernel, otherwise the existing host code. Nothing changes for host tiles.

Design

Each TA::Tensor member that has a device kernel wraps the device branch in TA_ON_DEVICE(types...) (device/dispatch.h). The macro expands to an if chain whose last else is the block written after it:

flowchart TD
    A["tile op, e.g. add_to"] --> B{"device_executable_v of the operand types<br/>(if constexpr)"}
    B -- "false: other tile types, CPU-only builds<br/>(discarded at compile time)" --> H["existing host code"]
    B -- true --> C{"execution_space() == Device<br/>(runtime)"}
    C -- "false: inside ExecutionSpaceScope(Host)" --> S["device::synchronize():<br/>wait for kernels this task already queued"] --> H
    C -- true --> D["device kernel<br/>(blaspp, thrust, librett)"] --> R["return result"]

Loading

For example, Tensor::add_to (tensor/tensor.h), after its existing early exits for empty tiles:

TA_ON_DEVICE(Tensor, Right) {
  TA_ASSERT(detail::is_range_set_congruent(*this, right) &&
            nbatch() == right.nbatch());
  device::axpy(total_size(), numeric_type(1), right.data(), data(),
               range_);
  return *this;
}
return inplace_binary(right, [](value_type& MADNESS_RESTRICT l,
                                const value_t<Right> r) { l += r; });

• Dispatch. The device branches are discarded at compile time for every other tile type and in builds without device support, so TA::Tensor needs no #ifdefs. Which tile types qualify is decided by the allocator's traits: allocator_memory_space_v<A> (platform.h) gives the memory space an allocator type allocates in (Device_UM for device_um_allocator), and device_executable_v requires a device-visible space, a BLAS element type, and the same type for every operand. is_device_tile<UMTensor<T>> is defined next to it, so every translation unit that uses TA::Tensor sees it.

• Kernels. gemm, copy, scal, axpy, mult, mult_to, permute, dot and inner_product are declared in device/dispatch.h and compiled once in device/tensor.cpp (blaspp, thrust, librett), so user translation units do not include device headers.

• Members with device branches. gemm, scaling, negation, add/subt (with factor and permutation), Hadamard product, permute, clone (also used by shift), the in-place forms, dot, inner_product and squared_norm. Ops with no device kernel (e.g. conj) synchronize the task's device work, then run the host code.

• ExecutionSpaceScope (platform.h) switches UMTensor ops to the host code for its lifetime, for debugging and as a reference. The setting is process-wide and read when each task runs; open scopes from the main thread only, after TA::initialize (the constructor throws otherwise). Under a Host scope, UnaryEvalImpl/BinaryEvalImpl skip the device task for these tiles. A task already running when a scope opens may have queued kernels before the switch; the Host branch of TA_ON_DEVICE waits for them before the host code reads the data.

  using Array = TA::DistArray<TA::UMTensor<double>>;
  Array c, c_host;
  c("i,j") = a("i,k") * b("k,j");  // tile ops run on the GPU
  {
    TA::ExecutionSpaceScope host(TA::ExecutionSpace::Host);
    c_host("i,j") = a("i,k") * b("k,j");  // same UMTensor arrays, host code
    world.gop.fence();  // finish the work before the scope ends
  }

• Expression engine. is_device_tile<UMTensor<T>> routes tile ops through madness::add_device_task, as for the existing BTAS device tile.

• Helpers (device/tensor.h): tile and array to_device/to_host prefetch, ta_tensor_to_um_tensor/um_tensor_to_ta_tensor (keeping every batch; a no-op when both tile types are the same, as in code generic over the tile type), and the tile and array conversions. Archiving needs no specialization: Tensor::serialize of a UMTensor waits for the task's device work and prefetches the data to the host before storing it, in every translation unit.

Changes that affect existing code

• sync_madness_task_with outside a device task now blocks until the stream finishes (external/device.h). Before, the thread-local stream was preset, so a call from a plain task (e.g. foreach, truncate) recorded nothing and returned while the kernel was still running. This applies to every device tile, including btasUMTensor. Outside device tasks, stream_for/blasqueue_for now pick the stream by tile offset instead of always using stream 0.
• PaRSEC's own GPU support is off by default (cmake/modules/FindOrFetchMADWorld.cmake). With it on, each PaRSEC context reserves 95% of each visible GPU's free memory at startup, and TA does not run PaRSEC device tasks. These are cache defaults, so a parent project that sets PARSEC_GPU_WITH_CUDA (or _HIP, _LEVEL_ZERO) first keeps its value.
• Storage freed inside a device task is released after the task's device work completes (external/device.h, device/device_task_fn.h). A UMTensor temporary (e.g. a lazy block argument materialized by BinaryWrapper, or the unpermuted sum in add(right, perm)) can drop its last reference while kernels reading it are still queued; returning it to Umpire's pool then would let another thread reuse the block under those kernels. Tensor's deleter now hands the release to the enclosing device task (device::defer_until_task_completes), which runs it when the task is destroyed, after its stream callback. Outside device tasks every kernel has already waited for its stream, so storage is released at once. This replaces the blocking device::synchronize() calls that the permuted add/subt/mult and axpy_to(perm) used to make before freeing their temporaries.
• Host reductions on UM tiles synchronize first (Tensor::reduce): a reduction task converts lazy tiles inside a device task, so a block's shift could queue a device copy that the host reduction read before it finished.
• MADNESS pin bumped to 3ff5ac06d (the merge of world: let several foreign threads schedule on the PaRSEC backend m-a-d-n-e-s-s/madness#849; the range from the previous pin also brings in unrelated worlddc/cloud archive-chunking and mra commits, none of which TA uses): with the PaRSEC backend, several threads that PaRSEC did not create, such as the thread that runs CUDA stream callbacks, can now submit tasks. Before, the second such thread aborted at 2 or more ranks.
• detail::gemm computes ldc once, used by the host, device and logging paths.
• TA_ON_DEVICE ends in its own else, so following it with another else is a compile error rather than silently dead host code.

Limits

• Non-ToT tiles only.
• Only device_um_allocator maps to a device memory space today; a TA::Tensor with another allocator runs the host code. A future device-only allocator would also need its host branch disabled, since the host code reads the same memory.
• Nested (tensor-of-tensor) and arena tiles run on the host; device support for them is follow-up work.
• BLAS-1 kernels run in chunks of at most INT_MAX elements (blaspp's length type is int); device permutes above INT_MAX elements throw TiledArray::Exception (librett uses 32-bit indices).
• Under a Host scope only the elementwise evaluators skip the device task; contraction, reduction, assignment-time conversion and clone tasks are still device tasks whose tile ops run the host code (documented on ExecutionSpaceScope).
• Batched tiles (nbatch() > 1): every device branch covers all batches. The host code does not, which predates this PR: the value-returning ops and clone (built on Tensor(other, op) and friends) return one batch, scale_to/neg_to touch only the first, and permute asserts nbatch() == 1; only the binary in-place ops cover every batch.
• HIP builds have not been tested.

Follow-up

Once this PR is merged, a follow-up PR will remove the BTAS-based device tile (btasUMTensor), which UMTensor replaces.

Tests and examples

• tests/expressions_device_tensor.cpp: 79 cases comparing UMTensor expressions with the same expressions on host TA::Tensor arrays, and, for complex tiles and batched products, with the same UMTensor arrays under a Host scope. Also covers conversions, archives, the stream choice outside device tasks, the reduction race, a Host scope opened while device tasks run, tile ops on batched (nbatch() > 1) tiles, SparsePolicy arrays (shape estimation, contraction, norm, truncate), and release of storage freed inside a device task only after its stream drains.
• device_tensor_large_suite (disabled by default, about 18 GB per case): the INT_MAX chunking and the permute exception.
• Examples: examples/device/ta_dense_um_tensor.cpp, ta_vector_um_tensor.cpp.

New public API

TA::UMTensor<T> (fwd.h); ExecutionSpace, execution_space(), ExecutionSpaceScope and allocator_memory_space_v (platform.h); TA_ON_DEVICE and device_executable_v (device/dispatch.h). UMTensor is now an is_device_tile.

Testing

Check	Result
Full ta_test, CUDA (A30), 1 rank	2305/2309 passed, 4 skipped
Full ta_test, CUDA (A30), 2 ranks	all passed except an intermittent tolerance miss in the existing um_expressions_suite/dot (BTAS tiles, unchanged here)
Device suite, 1 rank, 1 and 3 streams; 2 ranks	76/76
After the review fixes (deferred release, batched tiles, serialize), GTX 1080, CUDA 12.9, Debug: device suite, 1 rank, 1 and 3 streams	79/79
After the review fixes: ta_test built from the device, expressions, sparse expressions and range units, 1 rank / 2 ranks	all passed
Large-tile suite	passes
CPU-only build, single-rank suite	2132/2134 passed, 2 skipped
MPQC device validation (CCSD(T), CCk GS/LR/EOM, spin-orbital variants), Pthreads backend	np1: 7/7; np2: 6/6
MPQC device validation, PaRSEC backend, MADNESS 3ff5ac06d	np1: 7/7; np2: 6/7 (the new h2o-ccsd-631g-pvdz-gpu test crashes at exit in MPQC's teardown, after the result)
Dense GEMM throughput, one A100, FP64	19.25 TFLOPS (98.7% of peak, 99.2% of cuBLAS)
CCSD PPL term, one A100, FP64	17.57 TFLOPS (90% of peak); same as the BTAS UM tile

[evaleev]

@ajay-mk if we can ever hope to be able to use this for CSV-CC we should probably avoid duplicating TA::Tensor tile ops (gemm, etc.); these are now pretty complex due to the need to manage arena memory. TA::Tensor is already allocator aware, so we can compose UM-based variant. However, we need to be able to control dispatch to host/device execution spaces based on the allocator's traits. I suggest we control this dynamically (i.e. not "if Allocator allocates in a device-visible memory we launch device kernels", but using a flag; this will also be useful for testing).

To minimize the amount of boilerplate in the tile ops it should be possible to use macros.

[ajay-mk]

I like the idea. I think the runtime flag can also give us control the execution space of specific operations, so that we can push only certain terms/operations into device. I will look into this.
I did test the current version with CCk, it runs, and results match the host side results for GS, LR and EOM.