mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
26a91a9f04
Summary: Modified files in `benchmarks/tensorexpr` to add support for NVIDIA's Fuser for the jit compiler. This support has some modifications besides adding an option to support the NVIDIA fuser: * Adds FP16 Datatype support * Fixes SOL/Algo calculations to generally use the data type instead of being fixed to 4 bytes * Adds IR printing and kernel printing knobs * Adds a knob `input_iter` to create ranges of inputs currently only for reductions * Adds further reduction support for Inner and Outer dimension reductions that are compatible with the `input_iter` knob. * Added `simple_element`, `reduce2d_inner`, and `reduce2d_outer` to isolate performance on elementwise and reduction operations in the most minimal fashion. Pull Request resolved: https://github.com/pytorch/pytorch/pull/44101 Reviewed By: ngimel Differential Revision: D23713658 Pulled By: bertmaher fbshipit-source-id: d6b83cfab559aefe107c23b3c0f2df9923b3adc1
82 lines
2.1 KiB
Python
82 lines
2.1 KiB
Python
from . import benchmark
|
|
from . import tensor_engine
|
|
|
|
|
|
class NormalizationBench(benchmark.Benchmark):
|
|
def __init__(self, mode, device, dtype, N, C, H, W):
|
|
super().__init__(mode, device, dtype)
|
|
self.N = N
|
|
self.C = C
|
|
self.H = H
|
|
self.W = W
|
|
|
|
self.data = self.nchw_rand(
|
|
[self.N, self.C, self.H, self.W],
|
|
device=device, dtype=dtype,
|
|
requires_grad=self.requires_grad,
|
|
)
|
|
self.running_mean = self.rand([self.C], device=device, dtype=dtype)
|
|
self.running_var = self.rand([self.C], device=device, dtype=dtype)
|
|
self.training = self.mode == "both"
|
|
|
|
def config(self):
|
|
return [self.N, self.C, self.H, self.W]
|
|
|
|
def memory_workload(self):
|
|
if self.mode == "fwd":
|
|
sol_count = 1 + 1
|
|
algorithmic_count = 2 + 1
|
|
else:
|
|
sol_count = (1 + 1) + (1 + 1)
|
|
algorithmic_count = (2 + 1) + (3 + 1)
|
|
|
|
buffer_size = self.N * self.C * self.H * self.W * 4
|
|
return {
|
|
"sol": buffer_size * sol_count,
|
|
"algorithmic": buffer_size * algorithmic_count,
|
|
}
|
|
|
|
@staticmethod
|
|
def default_configs():
|
|
return [[128, 32, 128, 128]]
|
|
|
|
|
|
class BatchNormBench(NormalizationBench):
|
|
def forward(self):
|
|
y = self.batch_norm(
|
|
self.data, self.running_mean, self.running_var, training=self.training
|
|
)
|
|
return y
|
|
|
|
@staticmethod
|
|
def module():
|
|
return "batchnorm"
|
|
|
|
|
|
class InstanceNormBench(NormalizationBench):
|
|
def forward(self):
|
|
y = self.instance_norm(self.data)
|
|
return y
|
|
|
|
@staticmethod
|
|
def module():
|
|
return "instance_norm"
|
|
|
|
def is_supported(self):
|
|
return tensor_engine.is_supported(self.instance_norm)
|
|
|
|
|
|
class LayerNormBench(NormalizationBench):
|
|
def forward(self):
|
|
y = self.layer_norm(self.data, [self.H, self.W])
|
|
return y
|
|
|
|
@staticmethod
|
|
def module():
|
|
return "layernorm"
|
|
|
|
|
|
benchmark.register_benchmark_class(BatchNormBench)
|
|
benchmark.register_benchmark_class(InstanceNormBench)
|
|
benchmark.register_benchmark_class(LayerNormBench)
|