add unit test for linear fuse relu

Author: zhuhaozhe

tests/cpu/test_linear_fuse_relu.py

@@ -0,0 +1,82 @@
+import torch
+import time
+import intel_pytorch_extension_py as ipex
+K=1 #128
+C=16 #64
+MB = 28
+
+class Cast(torch.nn.Module):
+    __constants__ = ['to_dtype']
+
+    def __init__(self, to_dtype):
+        super(Cast, self).__init__()
+        self.to_dtype = to_dtype
+
+    def forward(self, input):
+        return input.to(self.to_dtype)
+
+    def extra_repr(self):
+        return 'to(%s)' % self.to_dtype
+
+def get_rand_seed():
+    return int(time.time() * 1000000000)
+
+def _ipxex_linear_relu(random_seed, data_type = torch.float32):
+  torch.manual_seed(random_seed)
+  fc = ipex.LinearFuseRelu(C, K).to(data_type).to('dpcpp')
+  return fc
+
+def _cpu_linear_relu(random_seed, data_type = torch.float32):
+  torch.manual_seed(random_seed)
+  fc = torch.nn.ModuleList()
+  fc.append(torch.nn.Linear(C, K).to(data_type))
+  if data_type == torch.bfloat16:
+    fc.append(Cast(torch.float32))
+  fc.append(torch.nn.ReLU())
+  return torch.nn.Sequential(*fc)
+
+def _run_mlp(random_seed, fc_module, data_type = torch.float32, device='cpu'):
+  torch.manual_seed(random_seed)
+  x1 = torch.randn(MB, C).to(data_type).to(device).requires_grad_()
+  y1 = fc_module(x1)
+  z1 = y1.mean()
+  z1.backward()
+  if type(fc_module) == torch.nn.modules.container.Sequential:
+    return x1.grad, fc_module[0].weight.grad, fc_module[0].bias.grad
+  return x1.grad, fc_module.weight.grad, fc_module.bias.grad
+
+for data_type in [torch.float32, torch.bfloat16]:
+  seed = get_rand_seed()
+  ipex_fc = _ipxex_linear_relu(seed, data_type)
+  cpu_fc = _cpu_linear_relu(seed, data_type)
+
+  rtol = 1e-5
+  atol = rtol
+  if data_type == torch.bfloat16:
+    rtol = 1e-2
+    atol = rtol
+
+  seed = get_rand_seed()
+  input_grad_ipex, weight_grad_ipex, bias_grad_ipex = _run_mlp(seed, ipex_fc, data_type, device='dpcpp')
+  input_grad_cpu, weight_grad_cpu, bias_grad_cpu = _run_mlp(seed, cpu_fc, data_type)
+
+  if input_grad_ipex is None:
+    if input_grad_cpu is not None:
+      print("##################### {} linear fuse relu input grad FAIL".format(str(data_type)))
+    else:
+      print("##################### {} linear fuse relu input grad PASS".format(str(data_type)))
+  else:
+    if not input_grad_ipex.to(torch.float32).allclose(input_grad_cpu.to(torch.float32), rtol=rtol, atol=atol):
+      print("##################### {} linear fuse relu input grad FAIL".format(str(data_type)))
+    else:
+      print("##################### {} linear fuse relu input grad PASS".format(str(data_type)))
+
+  if not weight_grad_ipex.to(torch.float32).allclose(weight_grad_cpu.to(torch.float32), rtol=rtol, atol=atol):
+    print("##################### {} linear fuse relu weight grad FAIL".format(str(data_type)))
+  else:
+    print("##################### {} linear fuse relu weight grad PASS".format(str(data_type)))
+
+  if not bias_grad_ipex.to(torch.float32).allclose(bias_grad_cpu.to(torch.float32), rtol=rtol, atol=atol):
+    print("##################### {} linear fuse relu bias grad FAIL".format(str(data_type)))
+  else:
+    print("##################### {} linear fuse relu bias grad PASS".format(str(data_type)))