taco: add the stateful random data loader for TACO

Add similar logic as in python to load the same random tensor from disk.

Author: rohany

taco/bench.cpp

@@ -0,0 +1,31 @@
+#include "bench.h"
+
+#include <cstdlib>
+#include <iostream>
+
+#include "taco/tensor.h"
+#include "taco/util/strings.h"
+
+std::string constructRandomTensorKey(std::vector<int> dims, float sparsity) {
+  auto path = std::getenv("TACO_RANDOM_TENSOR_PATH");
+  if (path == nullptr) {
+    std::cout << "TACO_RANDOM_TENSOR_PATH is unset" << std::endl;
+    assert(false);
+  }
+  std::string pathStr(path);
+  std::stringstream result;
+  result << pathStr;
+  if (pathStr[pathStr.size() - 1] != '/') {
+    result << "/";
+  }
+  result << taco::util::join(dims, "x") << "-" << sparsity << ".tns";
+  return result.str();
+}
+
+taco::TensorBase loadRandomTensor(std::string name, std::vector<int> dims, float sparsity, taco::Format format) {
+  // For now, just say that the python code must generate the random
+  // tensor before use.
+  auto tensor = taco::read(constructRandomTensorKey(dims, sparsity), format, true);
+  tensor.setName(name);
+  return tensor;
+}

taco/bench.h

@@ -2,6 +2,7 @@
 #define TACO_BENCH_BENCH_H
 
 #include "benchmark/benchmark.h"
+#include "taco/tensor.h"
 
 // Register a benchmark with the following options:
 // * Millisecond output display
@@ -37,4 +38,6 @@
   ->ReportAggregatesOnly(true)                  \
   ->UseRealTime()
 
+taco::TensorBase loadRandomTensor(std::string name, std::vector<int> dims, float sparsity, taco::Format format);
+
 #endif //TACO_BENCH_BENCH_H

taco/windowing.cpp

@@ -17,24 +17,13 @@ static void applyBenchSizes(benchmark::internal::Benchmark* b) {
 static void bench_add_sparse_window(benchmark::State& state, const Format& f) {
   int dim = state.range(0);
   auto sparsity = 0.01;
-  Tensor<float> matrix("A", {dim, dim}, f);
-
-  srand(4357);
-  for (int i = 0; i < dim; i++) {
-    for (int j = 0; j < dim; j++) {
-      float rand_float = (float)rand()/(float)(RAND_MAX);
-      if (rand_float < sparsity) {
-        matrix.insert({i, j}, (float) ((int) (rand_float*3/sparsity)));
-      }
-    }
-  }
+  Tensor<double> matrix = loadRandomTensor("A", {dim, dim}, sparsity, f);
   matrix.pack();
 
-
   for (auto _ : state) {
     // Setup.
     state.PauseTiming();
-    Tensor<float> result("B", {dim-2, dim-2}, f);
+    Tensor<double> result("B", {dim-2, dim-2}, f);
     IndexVar i, j;
     result(i, j) = matrix(i(1, dim-1), j(1, dim-1)) + matrix(i(1, dim-1), j(1, dim-1));
     result.compile();