text: #include <gtest/gtest.h>
  #include <torch/csrc/jit/runtime/static/impl.h>
  #include <torch/torch.h>
  #include <thread>
  
  #include "test_utils.h"
  
  using namespace torch;
  using namespace torch::jit;
  using namespace torch::jit::test;
  
  TEST(CpuFusion, Simple) {
    const auto simple_script = R"JIT(
      def forward(self, a, b):
          return (a + b).relu().tanh()
    )JIT";
  
    Module m("module");
    m.define(simple_script);
  
    StaticModuleOptions opts; // start with the defaults.
    opts.enable_tensorexpr_fusion = true;
  
    auto input1 = at::randn({2, 3});
    auto input2 = at::ones({2, 3});
  
    auto smodule = StaticModule(m, /* is_frozen */ false, opts, {input1, input2});
    StaticRuntime runtime(smodule);
  
    // Test with sample inputs
    {
      auto actual = runtime({input1, input2}, {});
      auto expect = at::tanh(at::relu(input1 + input2));
      EXPECT_TRUE(at::allclose(expect, actual.toTensor()));
    }
  
    // Test with different inputs
    {
      auto new_input1 = at::randn({5, 14});
      auto new_input2 = at::randn({5, 14});
      auto actual = runtime({new_input1, new_input2}, {});
      auto expect = at::tanh(at::relu(new_input1 + new_input2));
      EXPECT_TRUE(at::allclose(expect, actual.toTensor()));
    }
  }
  
  TEST(CpuFusion, FallbackGraph) {
    const auto simple_script = R"JIT(
      def forward(self, a, b):
          return (a + b).relu().tanh()
    )JIT";
  
    Module m("module");
    m.define(simple_script);
  
    StaticModuleOptions opts; // start with the defaults.
    opts.enable_tensorexpr_fusion = true;
  
    auto sample_input1 = at::randn({2, 3});
    auto sample_input2 = at::ones({2, 3});
    auto smodule = StaticModule(
        m, /* is_frozen */ false, opts, {sample_input1, sample_input2});
  
    StaticRuntime runtime(smodule);
  
    // The sample inputs above were contiguous. Now, use a strided input
    // to trigger running the fallback graph.
    {
      auto input1 = at::narrow(at::randn({2, 6}), 1, 0, 3);
      auto input2 = at::ones({2, 3});
      auto expect = at::tanh(at::relu(input1 + input2));
      auto actual = runtime({input1, input2}, {});
      EXPECT_TRUE(at::allclose(expect, actual.toTensor()));
    }
  
    // Test with strided inputs of different size.
    {
      auto input1 = at::narrow(at::randn({10, 30}), 1, 0, 25);
      auto input2 = at::randn({10, 25});
      auto expect = at::tanh(at::relu(input1 + input2));
      auto actual = runtime({input1, input2}, {});
      EXPECT_TRUE(at::allclose(expect, actual.toTensor()));
    }
  }
  
  TEST(CpuFusion, ParallelRuntimes) {
    const auto simple_script = R"JIT(
      def forward(self, a, b):
          return (a + b).relu().tanh()
    )JIT";
  
    Module m("module");
    m.define(simple_script);
  
    StaticModuleOptions opts; // start with the defaults.
    opts.enable_tensorexpr_fusion = true;
  
    auto sample_input1 = at::randn({2, 3});
    auto sample_input2 = at::ones({2, 3});
    auto smodule = StaticModule(
        m, /* is_frozen */ false, opts, {sample_input1, sample_input2});
  
    constexpr size_t kNumThreads = 2;
    std::vector<std::vector<std::pair<int, int>>> all_inputs;
    for (size_t id = 0; id < kNumThreads; ++id) {
      std::vector<std::pair<int, int>> thread_input = {
          {id, id + 1},
          {id + 10, id + 11},
          {id + 20, id + 21},
          {id + 30, id + 31},
          {id + 40, id + 41},
          {id + 50, id + 51},
          {id + 60, id + 61},
          {id + 70, id + 71}};
      all_inputs.emplace_back(std::move(thread_input));
    }
  
    auto exec_runtime = [&](size_t tid) {
      const auto& inputs = all_inputs[tid];
      StaticRuntime runtime(smodule);
      for (const auto& inp : inputs) {
        auto a = at::randn({inp.first, inp.second});
        auto b = at::randn({inp.first, inp.second});
        auto expect = at::tanh(at::relu(a + b));
        auto actual = runtime({a, b}, {});
        EXPECT_TRUE(at::allclose(expect, actual.toTensor()));
      }
    };
  
    std::vector<std::thread> threads;
    for (size_t id = 0; id < kNumThreads; ++id) {
      threads.emplace_back(exec_runtime, id);
    }
  
    for (auto& t : threads) {
      t.join();
    }
  }
  