Smaller things to fix before enabling TE by default

.jenkins/pytorch/macos-test.sh

@@ -63,7 +63,7 @@ test_python_all() {
   # Increase default limit on open file handles from 256 to 1024
   ulimit -n 1024
 
-  python test/run_test.py --verbose --exclude test_jit_profiling test_jit_legacy test_jit_fuser_legacy test_jit_fuser_profiling test_jit_fuser_te test_tensorexpr --determine-from="$DETERMINE_FROM"
+  python test/run_test.py --verbose --exclude test_jit_simple test_jit_legacy test_jit_fuser_legacy --determine-from="$DETERMINE_FROM"
 
   assert_git_not_dirty
 }

.jenkins/pytorch/test.sh

@@ -143,8 +143,8 @@ test_python_nn() {
   assert_git_not_dirty
 }
 
-test_python_ge_config_profiling() {
-  time python test/run_test.py --include test_jit_profiling test_jit_fuser_profiling test_jit_fuser_te --verbose --determine-from="$DETERMINE_FROM"
+test_python_ge_config_simple() {
+  time python test/run_test.py --include test_jit_simple --verbose --determine-from="$DETERMINE_FROM"
   assert_git_not_dirty
 }
 
@@ -154,7 +154,7 @@ test_python_ge_config_legacy() {
 }
 
 test_python_all_except_nn() {
-  time python test/run_test.py --exclude test_nn test_jit_profiling test_jit_legacy test_jit_fuser_legacy test_jit_fuser_profiling test_jit_fuser_te test_tensorexpr --verbose --determine-from="$DETERMINE_FROM"
+  time python test/run_test.py --exclude test_nn test_jit_simple test_jit_legacy test_jit_fuser_legacy --verbose --determine-from="$DETERMINE_FROM"
   assert_git_not_dirty
 }
 

.jenkins/pytorch/win-test-helpers/test_python_all_except_nn.bat

@@ -1,3 +1,3 @@
 call %SCRIPT_HELPERS_DIR%\setup_pytorch_env.bat
-cd test && python run_test.py --exclude test_jit_profiling test_jit_legacy test_jit_fuser_legacy test_jit_fuser_profiling test_jit_fuser_te test_tensorexpr --verbose --determine-from="%1" && cd ..
+cd test && python run_test.py --exclude test_jit_legacy test_jit_fuser_legacy --verbose --determine-from="%1" && cd ..
 if ERRORLEVEL 1 exit /b 1

test/run_test.py

@@ -57,10 +57,9 @@
     'test_type_hints',
     'test_utils',
     'test_namedtuple_return_api',
-    'test_jit_profiling',
+    'test_jit_simple',
     'test_jit_legacy',
     'test_jit_fuser_legacy',
-    'test_jit_fuser_profiling',
     'test_tensorboard',
     'test_namedtensor',
     'test_type_promotion',

torch/csrc/jit/passes/tensorexpr_fuser.cpp

@@ -13,7 +13,7 @@
 namespace torch {
 namespace jit {
 
-static bool texpr_fuser_enabled_ = false;
+static bool texpr_fuser_enabled_ = true;
 void setTensorExprFuserEnabled(bool val) {
   texpr_fuser_enabled_ = val;
 }
@@ -266,9 +266,9 @@ std::pair<graph_node_list::iterator, bool> scanNode(
 }
 
 void fuseTensorExprs(std::shared_ptr<Graph>& graph) {
-  //if (!tensorExprFuserEnabled()) {
-  //  return;
-  //}
+  if (!tensorExprFuserEnabled()) {
+    return;
+  }
   GRAPH_DUMP("Before TExprFuser: ", graph);
 
   // Get rid of dead code so that we don't waste effort fusing it.

torch/csrc/jit/runtime/graph_executor.cpp

@@ -779,8 +779,7 @@ void runNondiffOptimization(
   // Fuse the dequant - op - quant patterns into quantized ops
   QuantFusion(graph);
 
-  //FuseGraph(graph, strict_fuser_check);
-  // strict_fuser_check is synomous with ProfilingExecutor on
+  // strict_fuser_check is synonymous with ProfilingExecutor on
   // if `strict_fuser_check` is set to `true`, run TE by default
   // otherwise fallback to the legacy executor and legacy fuser
   if (strict_fuser_check) {

torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp

@@ -39,7 +39,7 @@ static std::atomic<bool> executor_mode{true};
 static std::atomic<bool> profiling_mode{false};
 #else
 static std::atomic<bool> executor_mode{true};
-static std::atomic<bool> profiling_mode{false};
+static std::atomic<bool> profiling_mode{true};
 #endif
 
 static std::atomic<size_t> num_profiled_runs{1};

torch/csrc/jit/tensorexpr/kernel.cpp

@@ -397,17 +397,14 @@ Tensor* TensorExprKernel::computeFourOperand(
 Tensor* TensorExprKernel::computeValue(const torch::jit::Value* v) {
   switch (v->node()->kind()) {
     case aten::add: {
-      if (v->node()->inputs().size () > 2){
-      return computeTwoOperandWithAlpha(
-          "aten_add", v, [](const ExprHandle& lhs, const ExprHandle& rhs) {
-            return lhs + rhs;
-          });
-      }else{
-      return computeTwoOperand(
-          "aten_add", v, [](const ExprHandle& lhs, const ExprHandle& rhs) {
-            return lhs + rhs;
-          });
-      }
+      auto add_lambda = [](const ExprHandle& lhs, const ExprHandle& rhs) {
+        return lhs + rhs;
+      };
+      TORCH_INTERNAL_ASSERT(
+          v->node()->inputs().size() == 2 || v->node()->inputs().size() == 3);
+      return (v->node()->inputs().size() > 2)
+          ? computeTwoOperandWithAlpha("aten_add", v, add_lambda)
+          : computeTwoOperand("aten_add", v, add_lambda);
     } break;
 
     case aten::_cast_Float: {

torch/testing/_internal/common_utils.py

@@ -105,10 +105,10 @@ def prof_meth_call(*args, **kwargs):
 args, remaining = parser.parse_known_args()
 if args.ge_config == 'legacy':
     GRAPH_EXECUTOR = ProfilingMode.LEGACY
-elif args.ge_config == 'profiling':
-    GRAPH_EXECUTOR = ProfilingMode.PROFILING
-else:
+elif args.ge_config == 'simple':
     GRAPH_EXECUTOR = ProfilingMode.SIMPLE
+else:
+    GRAPH_EXECUTOR = ProfilingMode.PROFILING
 
 TEST_BAILOUTS = args.test_bailouts
 TEST_IN_SUBPROCESS = args.subprocess