Fix ToDlPack memory space classification for offset views

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)
 {

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 <typename Op>
+    __MATX_INLINE__ auto GetStorageBasePointer(const Op &op) noexcept {
+      using ptr_type = const typename Op::value_type *;
+      if constexpr (requires { op.GetStorage(); }) {
+        return static_cast<ptr_type>(op.GetStorage().data());
+      }
+      else {
+        return static_cast<ptr_type>(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<typename matx::remove_cvref_t<typename Op::value_type>*>(op.Data());
+        // Pass in the base pointer, not a potentially offset pointer returned by op.Data()
+        auto ptr_strip = const_cast<typename matx::remove_cvref_t<typename Op::value_type>*>(GetStorageBasePointer(op));
         auto kind = GetPointerKind(ptr_strip);
 
         // Try to get pointer from cuda

include/matx/core/tensor.h

@@ -786,6 +786,12 @@ class tensor_t : public detail::tensor_impl_t<T,RANK,Desc> {
    *
    * 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,26 +1508,40 @@ 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.
     t->data             = const_cast<void *>(static_cast<const void *>(this->Data()));
     t->device.device_id = 0;
 
     // Determine where this memory resides
-    void *data_ptr = const_cast<void *>(static_cast<const void *>(this->Data()));
+    // Pass in the base pointer, not a potentially offset pointer
+    // returned by this->Data()
+    void *data_ptr = const_cast<void *>(static_cast<const void *>(this->GetStorage().data()));
     auto kind = GetPointerKind(data_ptr);
     [[maybe_unused]] auto mem_res = cuPointerGetAttributes(sizeof(attr)/sizeof(attr[0]), attr, data, reinterpret_cast<CUdeviceptr>(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;
       }

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<TestType>({5, 10, 20});
+  (t = zeros<TestType>(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<const void *>(s.Data()),
+            static_cast<const void *>(t.GetStorage().data()));
+  print(s, 1, 1, 2);
+
+  // print() is also instantiated on the tensor_impl_t base class
+  const detail::tensor_impl_t<TestType, 3> &ti = t;
+  print(ti, 1, 1, 2);
+
+  MATX_EXIT_HANDLER();
+}