Skip to content

Device support for TA::Tensor#531

Draft
ajay-mk wants to merge 39 commits intomasterfrom
ajay/expt/device-tensor-support
Draft

Device support for TA::Tensor#531
ajay-mk wants to merge 39 commits intomasterfrom
ajay/expt/device-tensor-support

Conversation

@ajay-mk
Copy link
Member

@ajay-mk ajay-mk commented Aug 19, 2025 •
edited
Loading

Implements device support TA::Tensor

This PR implements device support for TA::Tensor. A new device tile type TA::UMTensor is introduced, and the existing device tile type based on btas::Tensor will be removed.

Changelist

  • Introduce TA::UMTensor, aka TA::Tensor with UM allocators
  • Implement kernels for UMTensor operations
  • Add UMTensor based tests
    • Basic tests
    • Expressions
    • Einsum
    • Librett
  • Add examples using UMTensor
  • Utility functions
  • Remove all btas::Tensor based device code

Notes

[will be added later]

ajay-mk added 30 commits July 21, 2025 10:19
@ajay-mk
docs: fix a bunch of typos
73bcba9
@ajay-mk
fwd: try introducing a TA::Tensor with UM allocator [skip ci]
595c613
@ajay-mk
fwd: TA::UMTensorType -> TA::UMTensor
285cb26 
UMTensor will be TA::Tensor type with a UM allocator

To avoid confusion, renames some existing uses of UMTensor
@ajay-mk
device: support operations on device for UMTensor
4613c8d 
Implemented everything except serialization support (needs testing)
@ajay-mk
device: move tile type independent array functions to a new file
c1af6ac
@ajay-mk
UMTensor: implement serialization in TA::Tensor style
70dd413
@ajay-mk
unit: rename btas um tensor tests
188513a
@ajay-mk
UMTensor: try serialization implementation BTASUMTensor style
cd9eee9
@ajay-mk
unit: introduce UMTensor tests
095203d 
Right now, its just a copy of the btas type test case
@ajay-mk
unit: try setting nbatch 1 for test cases
b52b007 
This forces to use the ctor with nbatches. Only nbatch == 1 is supported anyway.
@ajay-mk
UMTensor: instantiate UMTensor types
1369261
@ajay-mk
cuda.cmake: add forward-unknown-opts to CMAKE_CUDA_FLAGS
972eabe 
See https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html\#forward-unknown-opts-forward-unknown-opts
@ajay-mk
unit: rename UMTensor test fixture
4d7b6ee
@ajay-mk
device: dense multiplication example for UMTensor
f54725f
@ajay-mk
UMTensor: avoid circular dependency
d2084e6 
Don't include tensor.h in um_tensor.h
@ajay-mk
unit: add missing tensor header for TensorUM tests
eb87064
@ajay-mk
refactor: update templating of UMTensor functions, drop concepts, avo…
f7a76c3 
…ids ambiguous definitions
@ajay-mk
UMTensor: specially handle ComplexConjugate in scaling
d061305 
Follows the same logic in device/btas.h
@ajay-mk
UMTensor: fix ldc computing in gemm calls, add assertions
47bddfb
@ajay-mk
Tensor: rever accidental changes to tensor.h
4e1121e
@ajay-mk
UMTensor: cleanup
e9f7d95
@ajay-mk
UMTensor: introduce const and non-const versions of to_device
60d4abd
@ajay-mk
UMTensor: refactor serialization implementation
720f678
@ajay-mk
UMTensor: no need to use value_type, use T directly + format
e5b98d6
@ajay-mk
format: clang-format source files
1f400bb
@ajay-mk
UMTensor: implement is_device_tile
f106b02
@ajay-mk
unit: expression tests with UMTensor
e6c3ba5
@ajay-mk
UM: Add to_host functions for Tensor types. Also removes non-const ve…
739d264 
…rsion of to_device, since pre-fetching is not a mutating operation.
@ajay-mk
UMTensor: use generic device_data, no need for overload
880e654
@ajay-mk
UMTensor: restrict kernels to tensor of scalars
af551f1
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@ajay-mk