Fix single gpu expr scripts & Add environment variables for capturing CUDA graph

Author: ints81

expr/single_gpu/correctness_test/scripts/run_densenet_correct_test.sh

@@ -12,20 +12,24 @@ do
             echo "##########################################################"
             echo "  Batch size : ${BATCH}"
             echo "  Growth K : ${GROWTH_K}"
-	        echo "  Num iters : ${ITER}"
+	          echo "  Num iters : ${ITER}"
             echo "##########################################################"
             python /workspace/correctness_test/code/dense_base.py ${BATCH} ${GROWTH_K} ${ITER}
 
             export DO_OOO_BACKPROP="true"
             export OOO_CAPTURE_OP="cluster_1_1/xla_run"
             export OOO_CAPTURE_ITER=2
             export OOO_NUM_BLOCK_OVERLAP_FORWARD=88
+            export OOO_OVERLAP_START="B4"
+            export OOO_OVERLAP_END="B3"
             export OOO_USE_SUB_STREAM="true"
             python /workspace/correctness_test/code/dense_ooo.py ${BATCH} ${GROWTH_K} ${ITER}
             unset DO_OOO_BACKPROP
             unset OOO_CAPTURE_OP
             unset OOO_CAPTURE_ITER
             unset OOO_NUM_BLOCK_OVERLAP_FORWARD
+	          unset OOO_OVERLAP_START
+	          unset OOO_OVERLAP_END
             unset OOO_USE_SUB_STREAM
 
 	    python /workspace/correctness_test/code/logit_diff.py

expr/single_gpu/scripts/run_densenet_expr_ooo.sh

@@ -6,6 +6,8 @@ export DO_OOO_BACKPROP="true"
 export OOO_CAPTURE_OP="cluster_1_1/xla_run"
 export OOO_CAPTURE_ITER=3
 export OOO_NUM_BLOCK_OVERLAP_FORWARD=88
+export OOO_OVERLAP_START="B4"
+export OOO_OVERLAP_END="B3"
 export OOO_USE_SUB_STREAM="true"
 
 BATCH=$1
@@ -25,4 +27,6 @@ unset DO_OOO_BACKPROP
 unset OOO_CAPTURE_OP
 unset OOO_CAPTURE_ITER
 unset OOO_NUM_BLOCK_OVERLAP_FORWARD
+unset OOO_OVERLAP_START
+unset OOO_OVERLAP_END
 unset OOO_USE_SUB_STREAM

tensorflow/tensorflow/compiler/xla/service/gpu/gpu_executable.cc

@@ -101,6 +101,14 @@ GpuExecutable::GpuExecutable(
     if (!str_capture_iter.empty()) {
       capture_iter_ = std::stoi(str_capture_iter);
     }
+
+    const char* cstr_overlap_start_name = std::getenv("OOO_OVERLAP_START");
+    std::string str_overlap_start_name(cstr_overlap_start_name ? cstr_overlap_start_name : "NONE");
+    overlap_start_name_ = str_overlap_start_name;
+
+    const char* cstr_overlap_end_name = std::getenv("OOO_OVERLAP_END");
+    std::string str_overlap_end_name(cstr_overlap_end_name ? cstr_overlap_end_name : "NONE");
+    overlap_end_name_ = str_overlap_end_name;
   }
 }
 

tensorflow/tensorflow/compiler/xla/service/gpu/gpu_executable.h

@@ -208,23 +208,34 @@ class GpuExecutable : public Executable {
   std::string FORWARD_OVERLAP_GRAPH = "FORWARD_OVERLAP_WGRADS";
   std::string DEFAULT_GRAPH = "LAST_GRAPH";
 
+  std::string overlap_start_name_;
+  std::string overlap_end_name_;
+
   bool is_overlap_w_grad_op( std::string op_name, std::string hlo_name  ){
-        if ( op_name.find("B4") != std::string::npos && 
-             op_name.find("Conv2DBackpropFilter") != std::string::npos &&
-             op_name.find("Dummy") == std::string::npos &&
-             hlo_name.find("custom") != std::string::npos ) {
-            return true;
-        } else {
-            return false;
-        }
+    if (overlap_start_name_ == "NONE") {
+      return false;
+    }
+
+    if ( op_name.find(overlap_start_name_) != std::string::npos && 
+        op_name.find("Conv2DBackpropFilter") != std::string::npos &&
+        op_name.find("Dummy") == std::string::npos &&
+        hlo_name.find("custom") != std::string::npos ) {
+      return true;
+    } else {
+      return false;
+    }
   }
 
   bool is_remain_graph_start( std::string op_name  ){
-        if ( op_name.find("B3") != std::string::npos ) {
-            return true;
-        } else {
-            return false;
-        }
+    if (overlap_end_name_ == "NONE") {
+      return false;
+    }
+
+    if ( op_name.find(overlap_end_name_) != std::string::npos ) {
+      return true;
+    } else {
+      return false;
+    }
   }
 
   bool is_deleted_op( std::string op_name ){