code: #include <benchmark/benchmark.h>
  #include <torch/csrc/jit/tensorexpr/ir_simplifier.h>
  #include <torch/csrc/jit/tensorexpr/loopnest.h>
  #include <torch/csrc/jit/tensorexpr/tensor.h>
  #include <torch/csrc/jit/tensorexpr/llvm_codegen.h>
  #include <torch/torch.h>
  
  using namespace torch::jit::tensorexpr;
  
  static void log_sleef(benchmark::State& state) {
    KernelScope ks;
    auto N = VarHandle("N", kInt);
    Placeholder A("A", kFloat, {N});
    torch::jit::tensorexpr::Tensor* B =
        Compute("B", {N}, [&](const VarHandle& i) {
          return log(A.load(i));
        });
    LoopNest ln({B});
    ln.prepareForCodegen();
    ln.vectorizeInnerLoops();
    Stmt* s = ln.root_stmt();
    s = torch::jit::tensorexpr::IRSimplifier::simplify(s);
    std::vector<CodeGen::BufferArg> args;
    args.emplace_back(B);
    args.emplace_back(A);
    args.emplace_back(N);
    LLVMCodeGen cg(s, args);
    at::Tensor A_t = torch::abs(torch::randn({state.range(0)}));
    at::Tensor B_t = torch::randn({state.range(0)});
    auto B_ref = at::log(A_t);
    cg.call({B_t.data_ptr<float>(), A_t.data_ptr<float>(), state.range(0)});
    assert(at::allclose(B_t, B_ref));
    for (auto _ : state) {
      cg.call({B_t.data_ptr<float>(), A_t.data_ptr<float>(), state.range(0)});
    }
    state.counters["log/s"] = benchmark::Counter(
        uint64_t(state.range(0) * state.iterations()), benchmark::Counter::kIsRate);
  }
  
  static void log_fast(benchmark::State& state) {
    KernelScope ks;
    auto N = VarHandle("N", kInt);
    Placeholder A("A", kFloat, {N});
    torch::jit::tensorexpr::Tensor* B =
        Compute("B", {N}, [&](const VarHandle& i) {
          return fast_log(A.load(i));
        });
    LoopNest ln({B});
    ln.prepareForCodegen();
    ln.vectorizeInnerLoops();
    Stmt* s = ln.root_stmt();
    s = torch::jit::tensorexpr::IRSimplifier::simplify(s);
    std::vector<CodeGen::BufferArg> args;
    args.emplace_back(B);
    args.emplace_back(A);
    args.emplace_back(N);
    LLVMCodeGen cg(s, args);
    at::Tensor A_t = torch::abs(torch::randn({state.range(0)}));
    at::Tensor B_t = torch::randn({state.range(0)});
    auto B_ref = at::log(A_t);
    cg.call({B_t.data_ptr<float>(), A_t.data_ptr<float>(), state.range(0)});
    assert(at::allclose(B_t, B_ref));
    for (auto _ : state) {
      cg.call({B_t.data_ptr<float>(), A_t.data_ptr<float>(), state.range(0)});
    }
    state.counters["log/s"] = benchmark::Counter(
        uint64_t(state.range(0) * state.iterations()), benchmark::Counter::kIsRate);
  }
  
  static void log_aten(benchmark::State& state) {
    at::Tensor A_t = torch::abs(torch::randn({state.range(0)}));
    at::Tensor B_t = torch::randn({state.range(0)});
    for (auto _ : state) {
      at::native::log_out(B_t, A_t);
    }
    state.counters["log/s"] = benchmark::Counter(
        uint64_t(state.range(0) * state.iterations()), benchmark::Counter::kIsRate);
  }
  
  static void logit_fast(benchmark::State& state) {
    KernelScope ks;
    auto N = VarHandle("N", kInt);
    Placeholder A("A", kFloat, {N});
    torch::jit::tensorexpr::Tensor* B =
        Compute("B", {N}, [&](const VarHandle& i) {
            auto A_elem = A.load(i);
            return fast_log(A_elem / (FloatImm::make(1.0f) - A_elem));
        });
    LoopNest ln({B});
    ln.prepareForCodegen();
    ln.vectorizeInnerLoops();
    Stmt* s = ln.root_stmt();
    s = torch::jit::tensorexpr::IRSimplifier::simplify(s);
    std::vector<CodeGen::BufferArg> args;
    args.emplace_back(B);
    args.emplace_back(A);
    args.emplace_back(N);
    LLVMCodeGen cg(s, args);
    at::Tensor A_t = torch::abs(torch::randn({state.range(0)}));
    at::Tensor B_t = torch::randn({state.range(0)});
    auto B_ref = at::logit(A_t);
    cg.call({B_t.data_ptr<float>(), A_t.data_ptr<float>(), state.range(0)});
    assert(at::allclose(B_t, B_ref));
    for (auto _ : state) {
      cg.call({B_t.data_ptr<float>(), A_t.data_ptr<float>(), state.range(0)});
    }
    state.counters["logit/s"] = benchmark::Counter(
        uint64_t(state.range(0) * state.iterations()), benchmark::Counter::kIsRate);
  }
  
  static void logit_aten(benchmark::State& state) {
    at::Tensor A_t = torch::abs(torch::randn({state.range(0)}));
    at::Tensor B_t = torch::randn({state.range(0)});
    for (auto _ : state) {
      at::native::logit_out(B_t, A_t);
    }
    state.counters["logit/s"] = benchmark::Counter(
        uint64_t(state.range(0) * state.iterations()), benchmark::Counter::kIsRate);
  }
  
  BENCHMARK(log_sleef)
    ->Args({2<<5})
    ->Args({2<<8})
    ->Args({2<<12})
    ->Args({2<<14});
  BENCHMARK(log_fast)
    ->Args({2<<5})
    ->Args({2<<8})
    ->Args({2<<12})
    ->Args({2<<14});
  BENCHMARK(log_aten)
    ->Args({2<<5})
    ->Args({2<<8})
    ->Args({2<<12})
    ->Args({2<<14});
  BENCHMARK(logit_fast)
    ->Args({2<<5})
    ->Args({2<<8})
    ->Args({2<<12})
    ->Args({2<<14});
  BENCHMARK(logit_aten)
    ->Args({2<<5})
    ->Args({2<<8})
    ->Args({2<<12})
    ->Args({2<<14});
  