mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-08 07:39:33 +01: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
64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
import torch
|
|
|
|
|
|
class TorchTensorEngine(object):
|
|
def rand(self, shape, device=None, dtype=None, requires_grad=False):
|
|
return torch.rand(shape, device=device, dtype=dtype, requires_grad=requires_grad)
|
|
|
|
def randn(self, shape, device=None, dtype=None, requires_grad=False):
|
|
return torch.randn(shape, device=device, dtype=dtype, requires_grad=requires_grad)
|
|
|
|
def nchw_rand(self, shape, device=None, requires_grad=False):
|
|
return self.rand(shape, device=device, requires_grad=requires_grad)
|
|
|
|
def reset(self, _):
|
|
pass
|
|
|
|
def rand_like(self, v):
|
|
return torch.rand_like(v)
|
|
|
|
def numpy(self, t):
|
|
return t.cpu().numpy()
|
|
|
|
def mul(self, t1, t2):
|
|
return t1 * t2
|
|
|
|
def add(self, t1, t2):
|
|
return t1 + t2
|
|
|
|
def batch_norm(self, data, mean, var, training):
|
|
return torch.nn.functional.batch_norm(data, mean, var, training=training)
|
|
|
|
def instance_norm(self, data):
|
|
return torch.nn.functional.instance_norm(data)
|
|
|
|
def layer_norm(self, data, shape):
|
|
return torch.nn.functional.layer_norm(data, shape)
|
|
|
|
def sync_cuda(self):
|
|
torch.cuda.synchronize()
|
|
|
|
def backward(self, tensors, grad_tensors, _):
|
|
torch.autograd.backward(tensors, grad_tensors=grad_tensors)
|
|
|
|
def sum(self, data, dims):
|
|
return torch.sum(data, dims)
|
|
|
|
def softmax(self, data, dim=None):
|
|
return torch.nn.functional.softmax(data, dim)
|
|
|
|
def max_pool2d(self, data, kernel_size, stride=1):
|
|
return torch.nn.functional.max_pool2d(data, kernel_size, stride=stride)
|
|
|
|
def avg_pool2d(self, data, kernel_size, stride=1):
|
|
return torch.nn.functional.avg_pool2d(data, kernel_size, stride=stride)
|
|
|
|
def conv2d_layer(self, ic, oc, kernel_size, groups=1):
|
|
return torch.nn.Conv2d(ic, oc, kernel_size, groups=groups)
|
|
|
|
def matmul(self, t1, t2):
|
|
return torch.matmul(t1, t2)
|
|
|
|
def to_device(self, module, device):
|
|
return module.to(device)
|