diff --git a/include/matx/core/allocator.h b/include/matx/core/allocator.h index 7bdaf544d..c61546720 100644 --- a/include/matx/core/allocator.h +++ b/include/matx/core/allocator.h @@ -365,13 +365,16 @@ __MATX_INLINE__ bool IsAllocated(void *ptr) { * Get the kind of pointer based on an address * * Returns the memory kind of the pointer (device, host, managed, etc) based on - *a pointer address. This function should not be used in the data path since it - *takes a mutex and possibly loops through a std::map. Since Views can modify - *the address of the data pointer, the base pointer may not be what is passed in - * to this function, and therefore would not be in the map. However, finding the - *next lowest address that is in the map is a good enough approximation since we - *also offset in a positive direction from the base, and generally if you're in - *a specific address range the type of pointer is obvious anyways. + * a pointer address. This function should not be used in the data path since it + * takes a mutex. + * + * The lookup is an exact match against the address the allocator recorded, so the + * base pointer of the allocation must be passed in. Views can offset their data + * pointer into the allocation, and such a pointer is never found in the map, so + * MATX_INVALID_MEMORY is returned for it. + * + * Note: an earlier version of this comment described an approximate lookup that + * fell back to the next lowest recorded address. No such fallback exists here. **/ __MATX_INLINE__ matxMemorySpace_t GetPointerKind(void *ptr) { diff --git a/include/matx/core/print.h b/include/matx/core/print.h index ad3a9681b..6b944c039 100644 --- a/include/matx/core/print.h +++ b/include/matx/core/print.h @@ -502,6 +502,30 @@ namespace matx { } + /** + * @brief Return the base pointer of the storage backing a tensor, when available + * + * GetPointerKind() is an exact-address lookup into the allocation map, so it only ever + * resolves the base pointer the allocator recorded. A view's Data() may point partway + * into that allocation and would never be found there. Tensor types that carry their + * storage can hand back its base pointer; tensor_impl_t and dynamic_tensor_t only know + * Data(). Either way the caller must tolerate a miss, since the storage base is itself + * an offset pointer for a non-owning view such as RealView()/ImagView() + * + * @param op input Operator + * @return base pointer of the storage, or op.Data() if the type does not carry storage + */ + template + __MATX_INLINE__ auto GetStorageBasePointer(const Op &op) noexcept { + using ptr_type = const typename Op::value_type *; + if constexpr (requires { op.GetStorage(); }) { + return static_cast(op.GetStorage().data()); + } + else { + return static_cast(op.Data()); + } + } + /** * @brief Print a tensor's values to output file stream * @@ -567,7 +591,8 @@ namespace matx { // If the user is printing a tensor with a const pointer underlying the data, we need to do the lookup // as if it's not const. This is because the ownership decision is done at runtime instead of compile-time, // so even though the lookup will never be done, the compilation path happens. - auto ptr_strip = const_cast*>(op.Data()); + // Pass in the base pointer, not a potentially offset pointer returned by op.Data() + auto ptr_strip = const_cast*>(GetStorageBasePointer(op)); auto kind = GetPointerKind(ptr_strip); // Try to get pointer from cuda diff --git a/include/matx/core/tensor.h b/include/matx/core/tensor.h index d1de8ab9f..4fd1da51f 100644 --- a/include/matx/core/tensor.h +++ b/include/matx/core/tensor.h @@ -786,6 +786,12 @@ class tensor_t : public detail::tensor_impl_t { * * Only available on complex data types. * + * This view builds its own non-owning storage at the reinterpreted pointer + * rather than sharing the tensor's storage_, so GetStorage().data() on the + * result is that same offset pointer, not the original allocation's base. + * GetPointerKind()'s exact-match lookup misses on it, so classify it via + * the cuPointerGetAttributes fallback instead of the allocator map. + * * @returns tensor view of only real-valued components * */ @@ -828,17 +834,40 @@ MATX_LOOP_UNROLL /** * @brief Return the storage container from the tensor * + * Returns a copy, which shares ownership of the buffer with this tensor. Use + * this overload to hand storage to something that keeps it, such as + * make_tensor() or a sparse tensor constructor. + * * @return storage container */ __MATX_INLINE__ auto GetStorage() noexcept { return storage_; } + /** + * @brief Return the storage container from the tensor + * + * Returns a reference, so reading through it costs no reference count update. + * The buffer is kept alive by this tensor, not by the returned reference, so + * do not hold it past the lifetime of the tensor. + * + * @return const reference to the storage container + */ + __MATX_INLINE__ const auto &GetStorage() const noexcept { + return storage_; + } + /** * Create a view of only imaginary-valued components of a complex array * * Only available on complex data types. * + * This view builds its own non-owning storage at the reinterpreted pointer + * rather than sharing the tensor's storage_, so GetStorage().data() on the + * result is that same offset pointer, not the original allocation's base. + * GetPointerKind()'s exact-match lookup misses on it, so classify it via + * the cuPointerGetAttributes fallback instead of the allocator map. + * * @returns tensor view of only imaginary-valued components * */ @@ -1479,10 +1508,11 @@ MATX_LOOP_UNROLL auto *mt = new ManagedType; DLTensor *t = &mt->dl_tensor; - CUpointer_attribute attr[] = {CU_POINTER_ATTRIBUTE_MEMORY_TYPE, CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL}; + CUpointer_attribute attr[] = {CU_POINTER_ATTRIBUTE_MEMORY_TYPE, CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL, CU_POINTER_ATTRIBUTE_IS_MANAGED}; CUmemorytype mem_type; int dev_ord; - void *data[2] = {&mem_type, &dev_ord}; + int is_managed; + void *data[3] = {&mem_type, &dev_ord, &is_managed}; // DLPack carries mutability via flags (versioned API), not via pointer type. // Preserve const-export semantics by marking versioned tensors read-only below. @@ -1490,15 +1520,28 @@ MATX_LOOP_UNROLL t->device.device_id = 0; // Determine where this memory resides - void *data_ptr = const_cast(static_cast(this->Data())); + // Pass in the base pointer, not a potentially offset pointer + // returned by this->Data() + void *data_ptr = const_cast(static_cast(this->GetStorage().data())); auto kind = GetPointerKind(data_ptr); [[maybe_unused]] auto mem_res = cuPointerGetAttributes(sizeof(attr)/sizeof(attr[0]), attr, data, reinterpret_cast(data_ptr)); MATX_ASSERT_STR_EXP(mem_res, CUDA_SUCCESS, matxCudaError, "Error returned from cuPointerGetAttributes"); if (kind == MATX_INVALID_MEMORY) { - if (mem_type == CU_MEMORYTYPE_DEVICE) { + // GetStorage().data() is only guaranteed to be the true allocation base + // for storage that still shares its owning tensor's buffer (e.g. Slice()/ + // Permute()). A view with its own non-owning storage at an offset/ + // reinterpreted address (e.g. RealView()/ImagView()) lands here instead, + // so classify it from the driver's own record of that address. + // CU_POINTER_ATTRIBUTE_MEMORY_TYPE alone can't tell managed memory from + // pinned host memory, so check CU_POINTER_ATTRIBUTE_IS_MANAGED too. + if (mem_type == CU_MEMORYTYPE_DEVICE || is_managed) { t->device.device_type = kDLCUDA; t->device.device_id = dev_ord; } + else if (mem_type == CU_MEMORYTYPE_HOST) { + t->device.device_type = kDLCUDAHost; + t->device.device_id = dev_ord; + } else { t->device.device_type = kDLCPU; } diff --git a/test/00_operators/print_test.cu b/test/00_operators/print_test.cu index 7538e3c39..ac2b2df33 100644 --- a/test/00_operators/print_test.cu +++ b/test/00_operators/print_test.cu @@ -32,3 +32,38 @@ TYPED_TEST(OperatorTestsFloatAllExecs, Print) MATX_EXIT_HANDLER(); } + +// Compile coverage for the two shapes of operator PrintData() has to +// classify memory for: a view whose data pointer is offset into its allocation, +// and the tensor_impl_t base class, which carries no storage to take a base +// pointer from and so must fall back to Data(). The second case is a build +// regression guard. +// +// This asserts no printed output. Both spellings print correctly +// either way: an unresolved pointer falls through to cuPointerGetAttributes, +// which classifies offset pointers fine wherever the driver knows the +// allocation. The fix removes the dependency on that fallback. +TYPED_TEST(OperatorTestsFloatAllExecs, PrintOffsetViewAndTensorImpl) +{ + MATX_ENTER_HANDLER(); + using TestType = cuda::std::tuple_element_t<0, TypeParam>; + using ExecType = cuda::std::tuple_element_t<1, TypeParam>; + + ExecType exec{}; + + auto t = make_tensor({5, 10, 20}); + (t = zeros(t.Shape())).run(exec); + exec.sync(); + + // A slice starting partway into the allocation, so Data() is not the base pointer + auto s = t.Slice({1, 2, 3}, {4, 8, 15}); + ASSERT_NE(static_cast(s.Data()), + static_cast(t.GetStorage().data())); + print(s, 1, 1, 2); + + // print() is also instantiated on the tensor_impl_t base class + const detail::tensor_impl_t &ti = t; + print(ti, 1, 1, 2); + + MATX_EXIT_HANDLER(); +} diff --git a/test/00_tensor/DLPackTests.cu b/test/00_tensor/DLPackTests.cu index c76a24fc2..1d4bafab0 100644 --- a/test/00_tensor/DLPackTests.cu +++ b/test/00_tensor/DLPackTests.cu @@ -190,15 +190,123 @@ void SetVersionedManagedTensorShapeAndStrides( mt->dl_tensor.strides = ctx->strides; } +// Verifies that a view whose data pointer is offset from the start of its base +// allocation still exports the memory space of that allocation. GetPointerKind() +// is an exact-key lookup into the allocation map, so a view's offset pointer is +// never found there and ToDlPack must classify the memory from the base pointer +// held by the storage. Both the legacy and versioned paths are checked since +// they funnel through the same ToDlPackImpl. +template +void CheckOffsetViewDLPackExport(ViewType &v, const void *base, + DLDeviceType expected_device_type) +{ + using TestType = typename ViewType::value_type; + + auto *data_ptr = const_cast(static_cast(v.Data())); + ASSERT_NE(data_ptr, base) << "view must start at a non-zero offset into its allocation"; + + auto check_dl_tensor = [&](auto *dl) { + ASSERT_EQ(dl->dl_tensor.ndim, v.Rank()); + + // The view's offset pointer is exported as-is, with no additional byte offset + ASSERT_EQ(dl->dl_tensor.data, data_ptr); + ASSERT_EQ(dl->dl_tensor.byte_offset, 0U); + + ASSERT_EQ(dl->dl_tensor.device.device_type, expected_device_type); + + auto dlt = detail::TypeToDLPackType(); + ASSERT_EQ(dl->dl_tensor.dtype.code, dlt.code); + ASSERT_EQ(dl->dl_tensor.dtype.bits, dlt.bits); + ASSERT_EQ(dl->dl_tensor.dtype.lanes, dlt.lanes); + + for (int32_t r = 0; r < v.Rank(); r++) { + ASSERT_EQ(dl->dl_tensor.shape[r], v.Size(r)); + ASSERT_EQ(dl->dl_tensor.strides[r], v.Stride(r)); + } + }; + + // A view shares its storage with the tensor it came from, so compare the + // reference count against what it was before the export rather than to 1 + const auto ref_count = v.GetRefCount(); + + { + auto *dl = v.ToDlPack(); + ASSERT_EQ(v.GetRefCount(), ref_count + 1); + ASSERT_NO_FATAL_FAILURE(check_dl_tensor(dl)); + dl->deleter(dl); + ASSERT_EQ(v.GetRefCount(), ref_count); + } + + { + auto *dlv = v.ToDlPackVersioned(); + ASSERT_EQ(v.GetRefCount(), ref_count + 1); + ASSERT_NO_FATAL_FAILURE(check_dl_tensor(dlv)); + dlv->deleter(dlv); + ASSERT_EQ(v.GetRefCount(), ref_count); + } +} + +// Calls check(tensor, expected_device_type) once for a {5, 10, 20} tensor in +// each memory space a tensor can be exported from. Shared by the full-tensor +// and offset-view export tests so the list of spaces lives in one place +template +void ForEachExportableMemorySpace(CheckFunc &&check) +{ + { + // Default memory space -> exercises whatever the allocator handed back. On + // devices where cudaDevAttrConcurrentManagedAccess == 0 (e.g. Jetson), the + // allocator falls back to host pinned memory, so derive the expected device + // type from the base allocation + auto t = make_tensor({5, 10, 20}); + auto kind = GetPointerKind(t.GetStorage().data()); + + // Default allocation can only be MATX_MANAGED_MEMORY or MATX_HOST_MEMORY. + // Assert that explicitly so a future allocator change can't silently + // produce a wrong expectation below + ASSERT_TRUE(kind == MATX_MANAGED_MEMORY || kind == MATX_HOST_MEMORY); + ASSERT_NO_FATAL_FAILURE(check(t, kind == MATX_HOST_MEMORY ? kDLCUDAHost : kDLCUDA)); + } + + { + // Explicit host pinned memory. This is the case that separates a base + // pointer lookup from the cuPointerGetAttributes fallback: the fallback + // cannot tell pinned memory from plain host memory and reports kDLCPU + auto t = make_tensor({5, 10, 20}, MATX_HOST_MEMORY); + ASSERT_NO_FATAL_FAILURE(check(t, kDLCUDAHost)); + } + + { + // Explicit plain malloc host memory -> kDLCPU + auto t = make_tensor({5, 10, 20}, MATX_HOST_MALLOC_MEMORY); + ASSERT_NO_FATAL_FAILURE(check(t, kDLCPU)); + } + + { + // Explicit device memory -> kDLCUDA + auto t = make_tensor({5, 10, 20}, MATX_DEVICE_MEMORY); + ASSERT_NO_FATAL_FAILURE(check(t, kDLCUDA)); + } + + { + // Async device memory on the default stream -> kDLCUDA + auto t = make_tensor({5, 10, 20}, MATX_ASYNC_DEVICE_MEMORY); + ASSERT_NO_FATAL_FAILURE(check(t, kDLCUDA)); + } +} + template class DLPackTestsAll : public ::testing::Test { }; template class DLPackTestsFloatNonComplex : public ::testing::Test { }; +template +class DLPackTestsComplex : public ::testing::Test { +}; TYPED_TEST_SUITE(DLPackTestsAll, MatXAllTypesCUDAExec); TYPED_TEST_SUITE(DLPackTestsFloatNonComplex, MatXFloatNonComplexTypesCUDAExec); +TYPED_TEST_SUITE(DLPackTestsComplex, MatXComplexTypesCUDAExec); TYPED_TEST(DLPackTestsAll, ExportLegacyDLPack) { @@ -228,39 +336,109 @@ TYPED_TEST(DLPackTestsAll, ExportLegacyDLPack) ASSERT_EQ(t.GetRefCount(), 1); }; - { - // Default memory space -> exercises whatever the allocator handed - // back. On devices where cudaDevAttrConcurrentManagedAccess == 0 (e.g. - // Jetson), the allocator falls back to host pinned memory. This test - // derives the expected device type from GetPointerKind - auto t = make_tensor({5, 10, 20}); - auto kind = GetPointerKind(t.GetStorage().data()); - - // Default allocation can only be MATX_MANAGED_MEMORY or MATX_HOST_MEMORY. - // Assert that explicitly so a future allocator change can't silently - // produce a wrong expectation below - ASSERT_TRUE(kind == MATX_MANAGED_MEMORY || kind == MATX_HOST_MEMORY); - auto expected_device_type = (kind == MATX_HOST_MEMORY ? kDLCUDAHost : kDLCUDA); - ASSERT_NO_FATAL_FAILURE(check_dl_export(t, expected_device_type)); - } + ASSERT_NO_FATAL_FAILURE(ForEachExportableMemorySpace(check_dl_export)); - { - // Explicit device memory -> should resolve to kDLCUDA - auto t = make_tensor({5, 10, 20}, MATX_DEVICE_MEMORY); - ASSERT_NO_FATAL_FAILURE(check_dl_export(t, kDLCUDA)); - } + MATX_EXIT_HANDLER(); +} - { - // Explicit host pinned memory -> should resolve to kDLCUDAHost - auto t = make_tensor({5, 10, 20}, MATX_HOST_MEMORY); - ASSERT_NO_FATAL_FAILURE(check_dl_export(t, kDLCUDAHost)); - } +TYPED_TEST(DLPackTestsAll, ExportOffsetSlice) +{ + MATX_ENTER_HANDLER(); - { - // Explicit plain malloc host memory -> should resolve to kDLCPU - auto t = make_tensor({5, 10, 20}, MATX_HOST_MALLOC_MEMORY); - ASSERT_NO_FATAL_FAILURE(check_dl_export(t, kDLCPU)); - } + using TestType = cuda::std::tuple_element_t<0, TypeParam>; + + auto check_offset_slice = [](auto &t, DLDeviceType expected_device_type) { + const void *base = static_cast(t.GetStorage().data()); + + // Starts partway into the allocation, so Data() is no longer the base + // pointer the allocator recorded + auto s = t.Slice({1, 2, 3}, {4, 8, 15}); + + ASSERT_EQ(s.Size(0), 3); + ASSERT_EQ(s.Size(1), 6); + ASSERT_EQ(s.Size(2), 12); + + // Pin the absolute layout of the view. CheckOffsetViewDLPackExport compares + // the exported shape and strides against the view's own accessors, so + // without this a regression in Slice would move both sides together + ASSERT_EQ(s.Stride(0), 200); + ASSERT_EQ(s.Stride(1), 20); + ASSERT_EQ(s.Stride(2), 1); + ASSERT_EQ(s.Data() - static_cast(base), 1 * 200 + 2 * 20 + 3); + + ASSERT_NO_FATAL_FAILURE(CheckOffsetViewDLPackExport(s, base, expected_device_type)); + }; + + ASSERT_NO_FATAL_FAILURE(ForEachExportableMemorySpace(check_offset_slice)); + + MATX_EXIT_HANDLER(); +} + +TYPED_TEST(DLPackTestsAll, ExportPermutedOffsetSlice) +{ + MATX_ENTER_HANDLER(); + + using TestType = cuda::std::tuple_element_t<0, TypeParam>; + + auto check_permuted_offset_slice = [](auto &t, DLDeviceType expected_device_type) { + const void *base = static_cast(t.GetStorage().data()); + auto s = t.Slice({1, 2, 3}, {4, 8, 15}); + auto p = s.Permute({2, 0, 1}); + + // Permuting reorders sizes and strides but keeps the slice's offset pointer + ASSERT_EQ(p.Data(), s.Data()); + ASSERT_EQ(p.Size(0), s.Size(2)); + ASSERT_EQ(p.Size(1), s.Size(0)); + ASSERT_EQ(p.Size(2), s.Size(1)); + ASSERT_EQ(p.Stride(0), s.Stride(2)); + ASSERT_EQ(p.Stride(1), s.Stride(0)); + ASSERT_EQ(p.Stride(2), s.Stride(1)); + ASSERT_NO_FATAL_FAILURE(CheckOffsetViewDLPackExport(p, base, expected_device_type)); + }; + + ASSERT_NO_FATAL_FAILURE(ForEachExportableMemorySpace(check_permuted_offset_slice)); + + MATX_EXIT_HANDLER(); +} + +// ToDlPackImpl() needs to classify Real/ImagView via cuPointerGetAttributes +// +// Slicing first guarantees an offset for both RealView() and ImagView(), +// since RealView() alone on an un-sliced tensor keeps the same address as +// the base allocation and wouldn't exercise the fallback path at all. +TYPED_TEST(DLPackTestsComplex, ExportRealImagViewOfOffsetSlice) +{ + MATX_ENTER_HANDLER(); + + using TestType = cuda::std::tuple_element_t<0, TypeParam>; + + auto check_real_imag_view = [](auto &t, DLDeviceType expected_device_type) { + const void *base = static_cast(t.GetStorage().data()); + auto s = t.Slice({1, 2, 3}, {4, 8, 15}); + + auto rv = s.RealView(); + auto iv = s.ImagView(); + + ASSERT_EQ(static_cast(rv.GetStorage().data()), static_cast(rv.Data())); + ASSERT_EQ(static_cast(iv.GetStorage().data()), static_cast(iv.Data())); + ASSERT_NE(static_cast(rv.Data()), base); + ASSERT_NE(static_cast(iv.Data()), base); + + auto check_export = [&](auto &v) { + auto *dl = v.ToDlPack(); + ASSERT_EQ(dl->dl_tensor.device.device_type, expected_device_type); + dl->deleter(dl); + + auto *dlv = v.ToDlPackVersioned(); + ASSERT_EQ(dlv->dl_tensor.device.device_type, expected_device_type); + dlv->deleter(dlv); + }; + + ASSERT_NO_FATAL_FAILURE(check_export(rv)); + ASSERT_NO_FATAL_FAILURE(check_export(iv)); + }; + + ASSERT_NO_FATAL_FAILURE(ForEachExportableMemorySpace(check_real_imag_view)); MATX_EXIT_HANDLER(); }