*: add python logic to parametrize over FROSTT tensors

It'll be easier to write these python FROSTT tests than the C++ ones.

Author: rohany

Makefile

@@ -5,7 +5,7 @@ BENCHFLAGS := #"--benchmark-group-by=func"
 IGNORE += taco
 IGNORE_FLAGS := $(addprefix --ignore=, $(IGNORE)) 
 
-export TACO_RANDOM_TENSOR_PATH = data/random/
+export TACO_TENSOR_PATH = data/
 
 # To group benchmark output by benchmark, use BENCHFLAGS=--benchmark-group-by=func.
 # To additionally group by a parameterized value, add on ",param:<paramname>" to the

numpy/ufuncs.py

@@ -2,6 +2,7 @@
 from scipy.sparse import random, csr_matrix
 import sparse
 import pytest
+from util import TensorCollectionFROSTT
 
 # TODO (rohany): Ask hameer about this. pydata/sparse isn't happy when
 #  given this ufunc to evaluate.
@@ -85,3 +86,17 @@ def bench():
         return C
     tacoBench(bench)
     print("Result", bench())
+
+# Run benchmarks against the FROSTT collection.
+FROSTTTensors = TensorCollectionFROSTT()
+@pytest.mark.skip(reason="way too slow")
+@pytest.mark.parametrize("tensor", FROSTTTensors.getTensors(), ids=FROSTTTensors.getTensorNames())
+def bench_pydata_frostt_ufunc_sparse(tacoBench, tensor):
+    frTensor = tensor.load()
+    # TODO (rohany): Replace this with a call to the random generator.
+    other = sparse.random(frTensor.shape, density=0.001)
+    def bench():
+        print("Beginning ldexp")
+        c = numpy.ldexp(frTensor, other)
+        return c
+    tacoBench(bench)

numpy/util.py

@@ -1,6 +1,11 @@
 import scipy.sparse
 import sparse
 import os
+import glob
+
+# Get the path to the directory holding random tensors. Error out
+# if this isn't set.
+TENSOR_PATH = os.environ['TACO_TENSOR_PATH']
 
 # TnsFileLoader loads a tensor stored in .tns format.
 class TnsFileLoader:
@@ -87,12 +92,10 @@ def load(self, path):
 # sparsity. For example, a 250 by 250 tensor with sparsity 0.01
 # would have a key of 250x250-0.01.tns.
 def construct_random_tensor_key(shape, sparsity):
-    # Get the path to the directory holding random tensors. Error out
-    # if this isn't set.
-    path = os.environ['TACO_RANDOM_TENSOR_PATH']
+    path = TENSOR_PATH
     dims = "x".join([str(dim) for dim in shape])
     key = "{}-{}.tns".format(dims, sparsity)
-    return os.path.join(path, key)
+    return os.path.join(path, "random", key)
 
 # RandomPydataSparseTensorLoader should be used to generate
 # random pydata.sparse tensors. It caches the loaded tensors
@@ -134,3 +137,27 @@ def random(self, shape, sparsity):
             dok = scipy.sparse.dok_matrix(result)
             TnsFileDumper().dump_dict_to_file(shape, dict(dok.items()), key)
             return result
+
+# FROSTTTensor represents a tensor in the FROSTT dataset.
+class FROSTTTensor:
+    def __init__(self, path):
+        self.path = path
+
+    def __str__(self):
+        f = os.path.split(self.path)[1]
+        return f.replace(".tns", "")
+
+    def load(self):
+        return PydataSparseTensorLoader().load(self.path)
+
+# TensorCollectionFROSTT represents the set of all FROSTT tensors.
+class TensorCollectionFROSTT:
+    def __init__(self):
+        data = os.path.join(TENSOR_PATH, "FROSTT")
+        frostttensors = glob.glob(os.path.join(data, "*.tns"))
+        self.tensors = [FROSTTTensor(t) for t in frostttensors]
+
+    def getTensors(self):
+        return self.tensors
+    def getTensorNames(self):
+        return [str(tensor) for tensor in self.getTensors()]

taco/bench.cpp

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

taco/bench.h

@@ -13,7 +13,7 @@
   BENCHMARK(bench)                \
   ->Unit(benchmark::kMillisecond) \
   ->Repetitions(10)               \
-  ->Iterations(5)                 \
+  ->Iterations(1)                 \
   ->ReportAggregatesOnly(true)    \
   ->UseRealTime()