From aca351de9c1467222c1e3f99571766eee8643513 Mon Sep 17 00:00:00 2001
From: Rohan Yadav <rohany@alumni.cmu.edu>
Date: Sat, 13 Feb 2021 11:57:01 -0800
Subject: [PATCH] tensor: allow `compute` to bypass caching mechanism

This commit adds a boolean to the `TensorBase::compute` method allowing
for bypass of the compute caching mechanism. This is useful when
benchmarking the generated kernels.
---
 include/taco/tensor.h | 5 ++++-
 src/tensor.cpp        | 5 +++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/include/taco/tensor.h b/include/taco/tensor.h
index 25186c815..641b978b2 100644
--- a/include/taco/tensor.h
+++ b/include/taco/tensor.h
@@ -420,7 +420,10 @@ class TensorBase {
   void assemble();
 
   /// Compute the given expression and put the values in the tensor storage.
-  void compute();
+  /// If force is true, then tensor expression will be executed, whether or
+  /// or not the tensor expression has been evaluated already. force is useful
+  /// in situations like benchmarking.
+  void compute(bool force = false);
 
   /// Compile, assemble and compute as needed.
   void evaluate();
diff --git a/src/tensor.cpp b/src/tensor.cpp
index ce2e4190d..35148cfb2 100644
--- a/src/tensor.cpp
+++ b/src/tensor.cpp
@@ -766,9 +766,10 @@ void TensorBase::assemble() {
   }
 }
 
-void TensorBase::compute() {
+void TensorBase::compute(bool force) {
   taco_uassert(!needsCompile()) << error::compute_without_compile;
-  if (!needsCompute()) {
+  // Return if the tensor doesn't need compute and the computation isn't forced.
+  if (!needsCompute() && !force) {
     return;
   }
   setNeedsCompute(false);