Move Accessor to its own file and add tests

Author: samanklesaria

src/libtorchaudio/CMakeLists.txt

@@ -6,6 +6,7 @@ set(
   lfilter.cpp
   overdrive.cpp
   utils.cpp
+  accessor_tests.cpp
   )
 
 set(

src/libtorchaudio/accessor.h

@@ -0,0 +1,44 @@
+#pragma once
+
+#include <torch/torch.h>
+#include <type_traits>
+#include <cstdarg>
+
+template<unsigned int k, typename T, bool IsConst = true>
+class Accessor {
+  int64_t strides[k];
+  T *data;
+
+public:
+  using tensor_type = typename std::conditional<IsConst, const torch::Tensor&, torch::Tensor&>::type;
+
+  Accessor(tensor_type tensor) {
+    data = tensor.template data_ptr<T>();
+    for (int i = 0; i < k; i++) {
+      strides[i] = tensor.stride(i);
+    }
+  }
+
+  T index(...) {
+    va_list args;
+    va_start(args, k);
+    int64_t ix = 0;
+    for (int i = 0; i < k; i++) {
+        ix += strides[i] * va_arg(args, int);
+    }
+    va_end(args);
+    return data[ix];
+  }
+
+  template<bool C = IsConst>
+  typename std::enable_if<!C, void>::type set_index(T value, ...) {
+    va_list args;
+    va_start(args, value);
+    int64_t ix = 0;
+    for (int i = 0; i < k; i++) {
+        ix += strides[i] * va_arg(args, int);
+    }
+    va_end(args);
+    data[ix] = value;
+  }
+};

src/libtorchaudio/accessor_tests.cpp

@@ -0,0 +1,25 @@
+#include <libtorchaudio/accessor.h>
+#include <cstdint>
+#include <torch/torch.h>
+
+using namespace std;
+
+bool test_accessor(const torch::Tensor& tensor) {
+  int64_t* data_ptr = tensor.template data_ptr<int64_t>();
+  auto accessor = Accessor<3, int64_t>(tensor);
+  for (int i = 0; i < tensor.size(0); i++) {
+    for (int j = 0; j < tensor.size(1); j++) {
+      for (int k = 0; k < tensor.size(2); k++) {
+        auto check = *(data_ptr++) ==  accessor.index(i, j, k);
+        if (!check) {
+          return false;
+        }
+      }
+    }
+  }
+  return true;
+}
+
+TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
+  m.def("torchaudio::_test_accessor", &test_accessor);
+}

src/libtorchaudio/forced_align/cpu/compute.cpp

@@ -5,8 +5,7 @@
 #include <torch/csrc/stable/ops.h>
 #include <torch/csrc/inductor/aoti_torch/c/shim.h>
 #include <torch/csrc/inductor/aoti_torch/utils.h>
-#include <cstdarg>
-#include <type_traits>
+#include <libtorchaudio/accessor.h>
 
 
 using namespace std;
@@ -15,44 +14,7 @@ namespace torchaudio {
 namespace alignment {
 namespace cpu {
 
-template<unsigned int k, typename T, bool IsConst = true>
-class Accessor {
-  int64_t strides[k];
-  T *data;
-
-public:
-  using tensor_type = typename std::conditional<IsConst, const torch::Tensor&, torch::Tensor&>::type;
-  
-  Accessor(tensor_type tensor) {
-    data = tensor.template data_ptr<T>();
-    for (int i = 0; i < k; i++) {
-      strides[i] = tensor.stride(i);
-    }
-  }
 
-  T index(...) {
-    va_list args;
-    va_start(args, k);
-    int64_t ix = 0;
-    for (int i = 0; i < k; i++) {
-        ix += strides[i] * va_arg(args, int);
-    }
-    va_end(args);
-    return data[ix];
-  }
-
-  template<bool C = IsConst>
-  typename std::enable_if<!C, void>::type set_index(T value, ...) {
-    va_list args;
-    va_start(args, value);
-    int64_t ix = 0;
-    for (int i = 0; i < k; i++) {
-        ix += strides[i] * va_arg(args, int);
-    }
-    va_end(args);
-    data[ix] = value;
-  }
-};
 
 // Inspired from
 // https://github.com/flashlight/sequence/blob/main/flashlight/lib/sequence/criterion/cpu/ConnectionistTemporalClassificationCriterion.cpp

test/torchaudio_unittest/accessor_test.py

@@ -0,0 +1,7 @@
+import torch
+from torchaudio._extension import _IS_TORCHAUDIO_EXT_AVAILABLE
+
+if _IS_TORCHAUDIO_EXT_AVAILABLE:
+    def test_accessor():
+        tensor = torch.randint(1000, (5,4,3))
+        assert torch.ops.torchaudio._test_accessor(tensor)