Merge pull request #402 from Infinoid/extra-test-case-for-316

Add a SpTV+openmp+atomics test case for #316

test/tests-scheduling-eval.cpp

@@ -345,13 +345,6 @@ TEST(scheduling_eval, test_spmvCPU_temp) {
   A.pack();
 
 
-  IndexVar i0("i0"), i1("i1"), kpos("kpos"), kpos0("kpos0"), kpos1("kpos1");
-  TensorVar tj("tj", Float64);
-  IndexVar jw("iw");
-
-  y(i) = A(i, j) * x(j);
-  Access tjAccess = tj();
-
   y(i) = A(i, j) * x(j);
   IndexStmt stmt = y.getAssignment().concretize();
   stmt = stmt.parallelize(i, ParallelUnit::CPUThread, OutputRaceStrategy::Atomics);
@@ -370,6 +363,57 @@ TEST(scheduling_eval, test_spmvCPU_temp) {
   ASSERT_TENSOR_EQ(expected, y);
 }
 
+TEST(scheduling_eval, test_sptvCPU_temp) {
+  if (should_use_CUDA_codegen()) {
+    return;
+  }
+  int NUM_I = 1021/10;
+  int NUM_J = 1039/10;
+  int NUM_K = 1049/10;
+  float SPARSITY = .01;
+  Tensor<double> A("A", {NUM_I, NUM_J, NUM_K}, Format({Sparse, Sparse, Sparse}));
+  Tensor<double> x("x", {NUM_K}, Format({Dense}));
+  Tensor<double> y("y", {NUM_J}, Format({Dense}));
+
+  srand(4357);
+  for (int i = 0; i < NUM_I; i++) {
+    for (int j = 0; j < NUM_J; j++) {
+      for (int k = 0; k < NUM_K; k++) {
+        float rand_float = (float)rand()/(float)(RAND_MAX);
+        if (rand_float < SPARSITY) {
+          A.insert({i, j, k}, (double) ((int) (rand_float*3/SPARSITY)));
+        }
+      }
+    }
+  }
+
+  for (int k = 0; k < NUM_K; k++) {
+    float rand_float = (float)rand()/(float)(RAND_MAX);
+    x.insert({k}, (double) ((int) (rand_float*3/SPARSITY)));
+  }
+
+  x.pack();
+  A.pack();
+
+
+  y(j) = A(i, j, k) * x(k);
+  IndexStmt stmt = y.getAssignment().concretize();
+  stmt = stmt.reorder({i,j,k}).parallelize(j, ParallelUnit::CPUThread, OutputRaceStrategy::Atomics);
+
+  //printToFile("test_sptvCPU_temp", stmt);
+
+  y.compile(stmt);
+  y.assemble();
+  y.compute();
+
+  Tensor<double> expected("expected", {NUM_J}, Format({Dense}));
+  expected(j) = A(i, j, k) * x(k);
+  expected.compile();
+  expected.assemble();
+  expected.compute();
+  ASSERT_TENSOR_EQ(expected, y);
+}
+
 TEST(scheduling_eval, example_spmvCPU_splitpos) {
   if (should_use_CUDA_codegen()) {
     return;