Adding couple more APIs to KVTensorWrapper to bring partiy with torch…

…::Tensor (#3645) Summary: X-link: facebookresearch/FBGEMM#721 Differential Revision: D68934783
pytorch · Feb 3, 2025 · 807bf97 · 807bf97
1 parent 98d54f7
commit 807bf97
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/fbgemm_gpu/src/ssd_split_embeddings_cache/kv_tensor_wrapper.h b/fbgemm_gpu/src/ssd_split_embeddings_cache/kv_tensor_wrapper.h
@@ -47,7 +47,9 @@ class KVTensorWrapper : public torch::jit::CustomClassHolder {
       const int64_t length,
       const at::Tensor& weights);
 
-  c10::IntArrayRef size();
+  c10::IntArrayRef sizes();
+
+  c10::IntArrayRef strides();
 
   c10::ScalarType dtype();
 

diff --git a/fbgemm_gpu/src/ssd_split_embeddings_cache/kv_tensor_wrapper_cpu.cpp b/fbgemm_gpu/src/ssd_split_embeddings_cache/kv_tensor_wrapper_cpu.cpp
@@ -57,11 +57,16 @@ void KVTensorWrapper::set_range(
   FBEXCEPTION("Not implemented");
 }
 
-c10::IntArrayRef KVTensorWrapper::size() {
+c10::IntArrayRef KVTensorWrapper::sizes() {
   FBEXCEPTION("Not implemented");
   return shape_;
 }
 
+c10::IntArrayRef KVTensorWrapper::strides() {
+  FBEXCEPTION("Not implemented");
+  return shape_; // make linter happy.
+}
+
 c10::ScalarType KVTensorWrapper::dtype() {
   FBEXCEPTION("Not implemented");
   return options_.dtype().toScalarType();

diff --git a/fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_split_table_batched_embeddings.cpp b/fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_split_table_batched_embeddings.cpp
@@ -347,10 +347,20 @@ void KVTensorWrapper::set_range(
   }
 }
 
-c10::IntArrayRef KVTensorWrapper::size() {
+c10::IntArrayRef KVTensorWrapper::sizes() {
   return shape_;
 }
 
+c10::IntArrayRef KVTensorWrapper::strides() {
+  // Assume contiguous tensor.
+  std::vector<int64_t> strides(shape_.size(), 1);
+  for (int i = shape_.size() - 2; i > -1; i--) {
+    int prev = i + 1;
+    strides[i] = strides[prev] * std::max<int64_t>(shape_[prev], 1);
+  }
+  return strides;
+}
+
 c10::ScalarType KVTensorWrapper::dtype() {
   return options_.dtype().toScalarType();
 }
@@ -500,9 +510,10 @@ static auto kv_tensor_wrapper =
         .def_property("layout_str", &KVTensorWrapper::layout_str)
         .def_property(
             "shape",
-            &KVTensorWrapper::size,
+            &KVTensorWrapper::sizes,
             std::string(
-                "Returns the shape of the original tensor. Only the narrowed part is materialized."));
+                "Returns the shape of the original tensor. Only the narrowed part is materialized."))
+        .def_property("strides", &KVTensorWrapper::strides);
 
 TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
   m.def(