From b51121eb5b15d2f4deb2886b7cc0e9d86b57a988 Mon Sep 17 00:00:00 2001
From: Rohan Yadav <rohany@alumni.cmu.edu>
Date: Wed, 10 Mar 2021 17:27:09 -0800
Subject: [PATCH] taco: add a cache for tensors loaded from files

Relates to #14.
---
 taco/bench.h    | 63 +++++++++++++++++++++++++++++--------------------
 taco/ufuncs.cpp |  8 +++++--
 2 files changed, 43 insertions(+), 28 deletions(-)

diff --git a/taco/bench.h b/taco/bench.h
index 04cf3fc..c3a3162 100644
--- a/taco/bench.h
+++ b/taco/bench.h
@@ -38,9 +38,46 @@
   ->ReportAggregatesOnly(true)                  \
   ->UseRealTime()
 
+template<typename T>
+taco::Tensor<T> castToType(std::string name, taco::Tensor<double> tensor) {
+  taco::Tensor<T> result(name, tensor.getDimensions(), tensor.getFormat());
+  std::vector<int> coords(tensor.getOrder());
+  for (auto& value : taco::iterate<double>(tensor)) {
+    for (int i = 0; i < tensor.getOrder(); i++) {
+      coords[i] = value.first[i];
+    }
+    result.insert(coords, T(value.second));
+  }
+  result.pack();
+  return result;
+}
+
+struct TacoTensorFileCache {
+  template<typename T>
+  taco::Tensor<double> read(std::string path, T format) {
+    if (this->lastPath == path) {
+      return lastLoaded;
+    }
+    // TODO (rohany): Not worrying about whether the format was the same as what was asked for.
+    this->lastLoaded = taco::read(path, format);
+    this->lastPath = path;
+    return this->lastLoaded;
+  }
+
+  template<typename T, typename U>
+  taco::Tensor<T> readIntoType(std::string name, std::string path, U format) {
+    auto tensor = this->read<U>(path, format);
+    return castToType<T>(name, tensor);
+  }
+
+  taco::Tensor<double> lastLoaded;
+  std::string lastPath;
+};
+
 std::string getTacoTensorPath();
 taco::TensorBase loadRandomTensor(std::string name, std::vector<int> dims, float sparsity, taco::Format format);
 
+// TODO (rohany): Cache the tensor shifts too.
 template<typename T, typename T2>
 taco::Tensor<T> shiftLastMode(std::string name, taco::Tensor<T2> original) {
   taco::Tensor<T> result(name, original.getDimensions(), original.getFormat());
@@ -63,30 +100,4 @@ taco::Tensor<T> shiftLastMode(std::string name, taco::Tensor<T2> original) {
   return result;
 }
 
-template<typename T>
-taco::Tensor<T> castToType(std::string name, taco::Tensor<double> tensor) {
-  taco::Tensor<T> result(name, tensor.getDimensions(), tensor.getFormat());
-  std::vector<int> coords(tensor.getOrder());
-  for (auto& value : taco::iterate<double>(tensor)) {
-    for (int i = 0; i < tensor.getOrder(); i++) {
-      coords[i] = value.first[i];
-    }
-    result.insert(coords, T(value.second));
-  }
-  result.pack();
-  return result;
-}
-
-template<typename T>
-taco::Tensor<T> readIntoType(std::string name, std::string path, taco::ModeFormat format) {
-  auto tensor = taco::read(path, format);
-  return castToType<T>(name, tensor);
-}
-
-template<typename T>
-taco::Tensor<T> readIntoType(std::string name, std::string path, taco::Format format) {
-  auto tensor = taco::read(path, format);
-  return castToType<T>(name, tensor);
-}
-
 #endif //TACO_BENCH_BENCH_H
diff --git a/taco/ufuncs.cpp b/taco/ufuncs.cpp
index 20855b7..81bc545 100644
--- a/taco/ufuncs.cpp
+++ b/taco/ufuncs.cpp
@@ -154,6 +154,10 @@ static void applyBenchSizes(benchmark::internal::Benchmark* b) {
 TACO_BENCH_ARGS(bench_ufunc_sparse, xor_0.01, 0.01, "xor")->Apply(applyBenchSizes);
 TACO_BENCH_ARGS(bench_ufunc_sparse, rightShift_0.01, 0.01, ">>")->Apply(applyBenchSizes);
 
+// fileCache is a cache of paths to tensors so that reads of the same tensor
+// don't need to perform duplicate disk IO.
+TacoTensorFileCache fileCache;
+
 static void bench_frostt_ufunc(benchmark::State& state, std::string tnsPath, Func op) {
   auto path = getTacoTensorPath();
   auto frosttTensorPath = path;
@@ -164,7 +168,7 @@ static void bench_frostt_ufunc(benchmark::State& state, std::string tnsPath, Fun
   frosttTensorPath += tnsPath;
 
   // TODO (rohany): What format do we want to do here?
-  auto frosttTensor = readIntoType<int64_t>("frostt", frosttTensorPath, Sparse);
+  auto frosttTensor = fileCache.readIntoType<int64_t>("frostt", frosttTensorPath, Sparse);
   Tensor<int64_t> other = shiftLastMode<int64_t, int64_t>("other", frosttTensor);
 
   for (auto _ : state) {
@@ -239,7 +243,7 @@ static void bench_suitesparse_ufunc(benchmark::State& state, Func op) {
   auto tensorName = taco::util::split(filename, ".")[0];
   state.SetLabel(tensorName);
 
-  auto ssTensor = readIntoType<int64_t>("ssTensor", tensorPath, CSR);
+  auto ssTensor = fileCache.readIntoType<int64_t>("ssTensor", tensorPath, CSR);
   auto other = shiftLastMode<int64_t, int64_t>("other", ssTensor);
 
   state.counters["dimx"] = ssTensor.getDimension(0);