mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
29d759948e
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/66746 Modified loops in files under fbsource/fbcode/caffe2/ from the format `for(TYPE var=x0;var<x_max;x++)` to the format `for(const auto var: irange(xmax))` This was achieved by running r-barnes's loop upgrader script (D28874212) with some modification to exclude all files under /torch/jit and a number of reversions or unused variable suppression warnings added by hand. Test Plan: Sandcastle Reviewed By: malfet Differential Revision: D31705361 fbshipit-source-id: 33fd22eb03086d114e2c98e56703e8ec84460268
60 lines
1.3 KiB
C++
60 lines
1.3 KiB
C++
#include <benchmark/benchmark.h>
|
|
#include <c10/core/InferenceMode.h>
|
|
#include <c10/util/irange.h>
|
|
#include <torch/csrc/jit/codegen/fuser/interface.h>
|
|
#include <torch/torch.h>
|
|
|
|
using namespace torch::jit;
|
|
|
|
static const std::string two_adds = R"JIT(
|
|
def two_adds(self, x: Tensor, y: Tensor, z: Tensor) -> Tensor:
|
|
return x + y + z
|
|
)JIT";
|
|
|
|
static void FusedOverhead(benchmark::State& state) {
|
|
c10::InferenceMode mode;
|
|
overrideCanFuseOnCPU(true);
|
|
|
|
Module m("m");
|
|
m.define(two_adds);
|
|
|
|
auto x = torch::ones({1});
|
|
auto y = torch::ones({1});
|
|
auto z = torch::ones({1});
|
|
|
|
// Warmup.
|
|
for (const auto i : c10::irange(8)) {
|
|
(void)i; // Suppress unused variable warning
|
|
m.run_method("two_adds", x, y, z);
|
|
}
|
|
|
|
for (auto _ : state) {
|
|
m.run_method("two_adds", x, y, z);
|
|
}
|
|
}
|
|
|
|
static void UnfusedOverhead(benchmark::State& state) {
|
|
c10::InferenceMode guard;
|
|
overrideCanFuseOnCPU(false);
|
|
|
|
Module m("m");
|
|
m.define(two_adds);
|
|
|
|
auto x = torch::ones({1});
|
|
auto y = torch::ones({1});
|
|
auto z = torch::ones({1});
|
|
|
|
// Warmup.
|
|
for (const auto i : c10::irange(8)) {
|
|
(void)i; // Suppress unused variable warning
|
|
m.run_method("two_adds", x, y, z);
|
|
}
|
|
|
|
for (auto _ : state) {
|
|
m.run_method("two_adds", x, y, z);
|
|
}
|
|
}
|
|
|
|
BENCHMARK(FusedOverhead);
|
|
BENCHMARK(UnfusedOverhead);
|