taco: add partial ufunc benchmark for FROSTT tensors

Author: rohany

taco/bench.cpp

@@ -6,16 +6,20 @@
 #include "taco/tensor.h"
 #include "taco/util/strings.h"
 
-std::string constructRandomTensorKey(std::vector<int> dims, float sparsity) {
+std::string getTacoTensorPath() {
   auto path = std::getenv("TACO_TENSOR_PATH");
   if (path == nullptr) {
     std::cout << "TACO_TENSOR_PATH is unset" << std::endl;
     assert(false);
   }
-  std::string pathStr(path);
+  return std::string(path);
+}
+
+std::string constructRandomTensorKey(std::vector<int> dims, float sparsity) {
+  auto path = getTacoTensorPath();
   std::stringstream result;
-  result << pathStr;
-  if (pathStr[pathStr.size() - 1] != '/') {
+  result << path;
+  if (path[path.size() - 1] != '/') {
     result << "/";
   }
   result << "random/";

taco/bench.h

@@ -34,10 +34,11 @@
   BENCHMARK_CAPTURE(bench, name, __VA_ARGS__)   \
   ->Unit(benchmark::kMicrosecond)               \
   ->Repetitions(10)                             \
-  ->Iterations(5)                               \
+  ->Iterations(1)                               \
   ->ReportAggregatesOnly(true)                  \
   ->UseRealTime()
 
+std::string getTacoTensorPath();
 taco::TensorBase loadRandomTensor(std::string name, std::vector<int> dims, float sparsity, taco::Format format);
 
 template<typename T>

taco/ufuncs.cpp

@@ -150,3 +150,61 @@ 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);
+
+static void bench_frostt_ufunc(benchmark::State& state, std::string tnsPath) {
+  auto path = getTacoTensorPath();
+  auto frosttTensorPath = path;
+  if (frosttTensorPath[frosttTensorPath.size() - 1] != '/') {
+    frosttTensorPath += "/";
+  }
+  frosttTensorPath += "FROSTT/";
+  frosttTensorPath += tnsPath;
+
+  // TODO (rohany): What format do we want to do here?
+  auto frosttTensor = read(frosttTensorPath, Sparse);
+  Tensor<double> other = shiftLastMode<double>("other", frosttTensor);
+
+  // TODO (rohany): Parametrize by the ufunc as well.
+  // TODO (rohany): Certain ufuncs need for operands to be of a certain type.
+  //  ldexp for example requires the right hand side to be an integer (for python
+  //  at least). Not sure how we'll handle that for this.
+  Func ldExp("2^", Ldexp(), leftIncAlgebra());
+
+  for (auto _ : state) {
+    state.PauseTiming();
+    Tensor<double> result("result", frosttTensor.getDimensions(), frosttTensor.getFormat());
+    switch (frosttTensor.getOrder()) {
+      case 4: {
+        IndexVar i, j, k, l;
+        result(i, j, k, l) = ldExp(frosttTensor(i, j, k, l), other(i, j, k, l));
+        break;
+      }
+      case 5: {
+        IndexVar i, j, k, l, m;
+        result(i, j, k, l, m) = ldExp(frosttTensor(i, j, k, l, m), other(i, j, k, l, m));
+        break;
+      }
+      default:
+        state.SkipWithError("invalid tensor dimension");
+        return;
+    }
+    result.compile();
+    result.assemble();
+    state.ResumeTiming();
+
+    result.compute();
+  }
+}
+
+// TODO (rohany): We can define another macro to "nest" defining benchmarks
+//  for each of the ufuncs that we want to operate on.
+#define FOREACH_FROSTT_TENSOR(__func__) \
+  __func__(nips, "nips.tns") \
+  __func__(uber_pickups, "uber-pickups.tns") \
+  __func__(chicaco_crime, "chicago-crime.tns") \
+  __func__(lbnl_network, "lbnl-network.tns")
+
+#define DECLARE_FROSTT_UFUNC_BENCH(name, path) \
+  TACO_BENCH_ARGS(bench_frostt_ufunc, name, path);
+
+FOREACH_FROSTT_TENSOR(DECLARE_FROSTT_UFUNC_BENCH)