From f6392d0150672b6773cdf3acff619246559c21de Mon Sep 17 00:00:00 2001
From: Rohan Yadav <rohany@lanka.csail.mit.edu>
Date: Fri, 2 Apr 2021 15:24:48 -0400
Subject: [PATCH] taco,numpy: add a benchmark to see the scaling of fused
 ufuncs

---
 numpy/ufuncs.py               | 7 +++----
 scripts/frostt_taco_runner.sh | 6 ++++--
 taco/ufuncs.cpp               | 9 +++------
 taco/windowing.cpp            | 1 -
 4 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/numpy/ufuncs.py b/numpy/ufuncs.py
index 233af90..15fa531 100644
--- a/numpy/ufuncs.py
+++ b/numpy/ufuncs.py
@@ -34,14 +34,14 @@ def bench():
         C = ufunc(A, B)
     tacoBench(bench)
 
-@pytest.mark.parametrize("dim", [5000, 10000, 20000])
+@pytest.mark.parametrize("dim", [1000, 2500, 5000, 10000, 20000, 40000])
 def bench_pydata_ufunc_fused(tacoBench, dim):
     loader = RandomPydataSparseTensorLoader()
     matrix = safeCastPydataTensorToInts(loader.random((dim, dim), 0.01))
     matrix1 = safeCastPydataTensorToInts(loader.random((dim, dim), 0.01, variant=1))
     matrix2 = safeCastPydataTensorToInts(loader.random((dim, dim), 0.01, variant=2))
     def bench():
-        result = numpy.logical_xor(numpy.logical_xor(matrix, matrix1), matrix2)
+        result = numpy.logical_and(numpy.logical_xor(matrix, matrix1), matrix2)
         return result
     tacoBench(bench)
 
@@ -153,7 +153,7 @@ def bench():
 fusedFuncNames = [
         "xorAndFused", 
         "xorOrFused",
-        #"xorXorFused", 
+        "xorXorFused", 
 ]
 fusedFuncs = zip(fusedFuncs, fusedFuncNames)
 @pytest.mark.parametrize("tensor", FROSTTTensors.getTensors())
@@ -163,7 +163,6 @@ def bench_pydata_frostt_fused_ufunc_sparse(tacoBench, tensor, func):
     third = PydataTensorShifter().shiftLastMode(other)
     def bench():
         c = func[0](frTensor, other, third)
-        print("nnz = ", c.nnz)
         return c
     extra_info = dict()
     extra_info['tensor_str'] = str(tensor)
diff --git a/scripts/frostt_taco_runner.sh b/scripts/frostt_taco_runner.sh
index 5031771..b39d66c 100755
--- a/scripts/frostt_taco_runner.sh
+++ b/scripts/frostt_taco_runner.sh
@@ -4,11 +4,13 @@
 #SBATCH -p lanka-v3
 #SBATCH --exclusive
 
-out=frostt-ufunc-bench/taco/
+# out=frostt-ufunc-bench/taco/
+out=frostt-fused-ufunc-bench/taco/
 
 mkdir -p "$out"
 
 for tensor in "nips" "uber-pickups" "chicago-crime" "enron" "nell-2" "vast"; do
 	csvout="$out/result-$tensor.csv"
-	LANKA=ON TACO_OUT="$csvout" make -j8 taco-bench BENCHES="bench_frostt_ufunc/$tensor*"
+	# LANKA=ON TACO_OUT="$csvout" make -j8 taco-bench BENCHES="bench_frostt_ufunc/$tensor*"
+	LANKA=ON TACO_OUT="$csvout" make -j8 taco-bench BENCHES="bench_frostt_ufunc_fused/$tensor*"
 done
diff --git a/taco/ufuncs.cpp b/taco/ufuncs.cpp
index eb3e857..cc485e2 100644
--- a/taco/ufuncs.cpp
+++ b/taco/ufuncs.cpp
@@ -228,18 +228,16 @@ static void bench_ufunc_fused(benchmark::State& state, const Format& f) {
     state.PauseTiming();
     Tensor<int64_t> result("result", {dim, dim}, f);
     IndexVar i("i"), j("j");
-    result(i, j) = nestedXorOp(matrix(i, j), matrix1(i, j), matrix2(i, j));
-    result.setAssembleWhileCompute(false);
+    result(i, j) = andOp(xorOp(matrix(i, j), matrix1(i, j)), matrix2(i, j));
+    result.setAssembleWhileCompute(true);
     result.compile();
     state.ResumeTiming();
     result.assemble();
     result.compute();
-    result = result.removeExplicitZeros(result.getFormat());
-
   }
 }
  TACO_BENCH_ARGS(bench_ufunc_fused, csr, CSR)
-   ->ArgsProduct({{5000, 10000, 20000}});
+   ->ArgsProduct({{1000, 2500, 5000, 10000, 20000, 40000}});
 
 // UfuncInputCache is a cache for the input to ufunc benchmarks. These benchmarks
 // operate on a tensor loaded from disk and the same tensor shifted slightly. Since
@@ -467,7 +465,6 @@ static void bench_frostt_ufunc_fused(benchmark::State& state, std::string tnsPat
     state.ResumeTiming();
 
     result.compute();
-    
   }
 }
 
diff --git a/taco/windowing.cpp b/taco/windowing.cpp
index bd7f7da..1ab1f0f 100644
--- a/taco/windowing.cpp
+++ b/taco/windowing.cpp
@@ -26,7 +26,6 @@ enum WindowConfig {
   __func__(Constant, Constant) \
   __func__(ConstantFraction, ConstantFraction) \
   __func__(AlmostWhole, AlmostWhole) \
-  __func__(Whole, Whole) \
   __func__(NoWindowing, NoWindowing)
 
 Tensor<double> windowedTensorOp(Tensor<double> input1, Tensor<double> input2, int dim, WindowConfig config, float sparsity) {
