Fix building for CUDA 12.4 and for torch>=2.4 (#1297)

k2-fsa · Jul 24, 2024 · 9897a3a · 9897a3a
1 parent 07c00d1
commit 9897a3a
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 6 deletions.
diff --git a/k2/csrc/ragged_ops.cu b/k2/csrc/ragged_ops.cu
@@ -2515,12 +2515,18 @@ struct HashOutputIteratorDeref {  // this is what you get when you dereference
 
 template <typename T>
 struct HashOutputIterator {  // outputs just the index of the pair.
-  explicit HashOutputIterator(T *t) : t_(t) {}
-  __device__ __forceinline__ HashOutputIteratorDeref<T> operator[](
+  explicit __host__ __device__ __forceinline__ HashOutputIterator(T *t)
+      : t_(t) {}
+  __host__ __device__ __forceinline__ HashOutputIteratorDeref<T> operator[](
       int32_t idx) const {
     return HashOutputIteratorDeref<T>(t_ + idx);
   }
-  __device__ __forceinline__ HashOutputIterator operator+(size_t offset) {
+  __host__ __device__ __forceinline__ HashOutputIteratorDeref<T> operator*()
+      const {
+    return HashOutputIteratorDeref<T>(t_);
+  }
+  __host__ __device__ __forceinline__ HashOutputIterator
+  operator+(size_t offset) {
     return HashOutputIterator{t_ + offset};
   }
   T *t_;

diff --git a/k2/csrc/ragged_ops_inl.h b/k2/csrc/ragged_ops_inl.h
@@ -578,12 +578,18 @@ struct PairOutputIteratorDeref {  // this is what you get when you dereference
 
 template <typename T>
 struct PairOutputIterator {  // outputs just the index of the pair.
-  explicit PairOutputIterator(int32_t *i) : i_(i) {}
-  __device__ __forceinline__ PairOutputIteratorDeref<T> operator[](
+  explicit __host__ __device__ __forceinline__ PairOutputIterator(int32_t *i)
+      : i_(i) {}
+  __host__ __device__ __forceinline__ PairOutputIteratorDeref<T> operator[](
       int32_t idx) const {
     return PairOutputIteratorDeref<T>(i_ + idx);
   }
-  __device__ __forceinline__ PairOutputIterator operator+(int32_t offset) {
+  __host__ __device__ __forceinline__ PairOutputIteratorDeref<T> operator*()
+      const {
+    return PairOutputIteratorDeref<T>(i_);
+  }
+  __host__ __device__ __forceinline__ PairOutputIterator
+  operator+(int32_t offset) {
     return PairOutputIterator{i_ + offset};
   }
   int32_t *i_;

diff --git a/k2/python/csrc/torch.h b/k2/python/csrc/torch.h
@@ -30,6 +30,14 @@
 #include "k2/python/csrc/torch.h"
 #include "torch/extension.h"
 
+#if K2_TORCH_VERSION_MAJOR > 2 || \
+    (K2_TORCH_VERSION_MAJOR == 2 && K2_TORCH_VERSION_MINOR >= 4)
+// For torch >= 2.4.x
+// do nothing to fix the following error
+// error: class "pybind11::detail::type_caster<c10::ScalarType, void>" has
+// already been defined
+#else
+// For torch < 2.4
 namespace pybind11 {
 namespace detail {
 
@@ -71,6 +79,7 @@ struct type_caster<torch::ScalarType> {
 
 }  // namespace detail
 }  // namespace pybind11
+#endif
 
 namespace k2 {
 /* Transfer an object to a specific device.