make OneDNN pooling call inference path

Author: XiaobingSuper

torch_ipex/csrc/cpu/CustomOPs.h

@@ -1,6 +1,8 @@
 #pragma once
 
 #include "DevOPs.h"
+#include "dbl/Common.h"
+#include "dil/dil.hpp"
 #include "torch_ipex/csrc/aten_ipex_bridge.h"
 #include "torch_ipex/csrc/utils.h"
 #include <ATen/Tensor.h>
@@ -150,9 +152,11 @@ class NewMaxPool2dOp : public torch::autograd::Function<NewMaxPool2dOp> {
     try {
       if (torch_ipex::check_auto_dnnl() &&
           input.device().type() == c10::DeviceType::DPCPP) {
+        auto src_dil_type = torch_ipex::cpu::dbl::comm::try_gen_dil_tensor(input).get_data_type();
+        auto input_temp = (src_dil_type == dil::data_type::u8 || src_dil_type == dil::data_type::s8) ? input : input.contiguous();
+
         at::Tensor output = torch_ipex::cpu::AtenIpexCPUDev::dil_max_pooling(
-            input.is_contiguous() ? input : input.contiguous(), kernel_size,
-            stride, padding, dilation, ceil_mode);
+            input_temp, kernel_size, stride, padding, dilation, ceil_mode);
         return std::tuple<at::Tensor, at::Tensor>(output, output);
       }
     } catch (std::exception &e) {
@@ -368,10 +372,10 @@ class NewApaptiveAvgPoolingOp
 public:
   static at::Tensor _forward(at::Tensor input, at::IntArrayRef output_size) {
     try {
-      if (torch_ipex::check_auto_dnnl() &&
-          input.device().type() == c10::DeviceType::DPCPP) {
-        return torch_ipex::cpu::AtenIpexCPUDev::dil_adaptive_avg_pool2d(
-            input.is_contiguous() ? input : input.contiguous(), output_size);
+      if (torch_ipex::check_auto_dnnl() && input.device().type() == c10::DeviceType::DPCPP) {
+        auto src_dil_type = torch_ipex::cpu::dbl::comm::try_gen_dil_tensor(input).get_data_type();
+        auto input_temp = (src_dil_type == dil::data_type::u8 || src_dil_type == dil::data_type::s8) ? input : input.contiguous();
+        return torch_ipex::cpu::AtenIpexCPUDev::dil_adaptive_avg_pool2d(input_temp, output_size);
       }
     } catch (std::exception &e) {
 #if defined(_DEBUG)

torch_ipex/csrc/cpu/DevOPs.cpp

@@ -264,7 +264,11 @@ at::Tensor AtenIpexCPUDev::dil_convolution_overrideable(const at::Tensor & input
           return AtenIpexCPUDev::dil_deconvolution(input.is_contiguous() ? input : input.contiguous(), weight.is_contiguous() ? weight : weight.contiguous(), bias.defined() && !bias.is_contiguous() ? bias.contiguous() : bias, padding, output_padding, stride, dilation, groups);
         } else {
           // for int8 path, input always acbd format which is non-contiguous, .contiguous() will reorder to fp32
-          return AtenIpexCPUDev::dil_convolution(input, weight, bias, stride, padding, dilation, groups);
+          auto src_dil_type = dbl::comm::try_gen_dil_tensor(input).get_data_type();
+          auto input_temp = (src_dil_type == dil::data_type::u8 || src_dil_type == dil::data_type::s8) ? input : input.contiguous();
+          auto weight_dil_type = dbl::comm::try_gen_dil_tensor(weight).get_data_type();
+          auto weight_temp = weight_dil_type == dil::data_type::s8 ? weight : weight.contiguous();
+          return AtenIpexCPUDev::dil_convolution(input_temp, weight_temp, bias, stride, padding, dilation, groups);
         }
       }
     }
@@ -788,7 +792,6 @@ at::Tensor AtenIpexCPUDev::dil_linear(
     b = dbl::comm::try_gen_dil_tensor(bias);
   }
 
-  auto output_scale = dbl::comm::get_int8_scale(/*  uint8_used=false */);
   dil::tensor y = dbl::linear::linear_impl(x, w, b, output_scale);
 
   auto input_size = self.sizes();

torch_ipex/csrc/cpu/dbl/Pool.cpp

@@ -125,6 +125,11 @@ at::Tensor _dil_pooling(
   }
 
   dil::tensor y;
+  dil::prop_kind aprop_kind = dil::prop_kind::forward;
+  auto src_type = x.get_data_type();
+  if (dil::data_type::s8 == src_type || dil::data_type::u8 == src_type) {
+    aprop_kind = dil::prop_kind::forward_inference;
+  }
   dil::pooling_forward::compute(
       x,
       {output_sizes.cbegin(), output_sizes.cend()},
@@ -134,7 +139,7 @@ at::Tensor _dil_pooling(
       {padding_vec_l.cbegin(), padding_vec_l.cend()},
       {padding_vec_r.cbegin(), padding_vec_r.cend()},
       algo,
-      dil::prop_kind::forward);
+      aprop_kind);
 
   return gen_aten_tensor_by(std::move(y));
 }