Prepack conv weights (#31)

Currently adopt prepacking conv weights in module.to Pros: thread-safe more explicit, compared to implicitly prepack at runtime Cons: no input info. meaning that queried format might not be optimal. This is not manifest in conv but this could be a problem for linear weights. unable to re-pack conv weights if it's been reordered back to plain (maybe even not a con)

Author: pinzhenx

intel_pytorch_extension_py/ops/__init__.py

@@ -5,4 +5,4 @@
 from .reshape import *
 from .mlp import * 
 from .linear_fuse_relu import *
-
+from .module import *

intel_pytorch_extension_py/ops/module.py

@@ -0,0 +1,26 @@
+import torch
+import _torch_ipex as core
+
+
+orig_module_to = torch.nn.Module.to
+
+def module_to(self, *args, **kwargs):
+ def prepack(m):
+ if isinstance(m, torch.nn.Conv2d) or isinstance(m, torch.nn.Conv3d):
+ core.prepack_conv_weight(m.weight, m.padding, m.stride, m.dilation, m.groups)
+
+ def prepack_reccur(m):
+ prepack(m)
+ for _, sub_m in m.named_children():
+ prepack_reccur(sub_m)
+
+ m = orig_module_to(self, *args, **kwargs)
+
+ device = torch._C._nn._parse_to(*args, **kwargs)[0]
+ if device and device.type == 'dpcpp':
+ prepack_reccur(m)
+
+ return m
+
+
+torch.nn.Module.to = module_to

torch_ipex/csrc/aten_ipex_bridge.cpp

@@ -11,7 +11,6 @@
 
 #include "ipex_tensor_impl.h"
 #include "ipex_sparse_tensor_impl.h"
-#include "cpu/dbl/Common.h"
 #include "cpu/ShadeDataContext.h"
 #include "cpu/bf16/Converter.h"
 #include "utils.h"
@@ -105,6 +104,41 @@ void reorderDilTensorToPublic(const at::Tensor& ipexTensor) {
  }
 }
 
+void reorderDilTensorGeneric(const at::Tensor& ipexTensor, const dil::tensor::desc& dstDesc) {
+ // ipexTensor is not required to be a DIL tensor
+ dil::tensor src = cpu::dbl::comm::try_gen_dil_tensor(ipexTensor);
+ dil::tensor dst {dstDesc};
+ dst.feed_from(src);
+
+ cpu::ShadeDataContext *new_shade_data_context = cpu::ShadeDataContext::allocShadeDataContext();
+ new_shade_data_context->data_type = cpu::SHADE_DATA_TYPE::DIL;
+ new_shade_data_context->dil_tensor = dst;
+
+ if (dstDesc.is_plain()) {
+ // Share with DNNL raw data because it is plain format now
+ new_shade_data_context->cpu_raw_data = dst.get_data_handle();
+ // Cannot free CPU data because the the data is owned by DNNL
+ new_shade_data_context->cpu_del_fun = &(c10::detail::deleteNothing);
+ } else {
+ // If tensor is of blocked format, cpu raw data means nothing here.
+ new_shade_data_context->cpu_raw_data = nullptr;
+ new_shade_data_context->cpu_del_fun = nullptr;
+ }
+
+ // Create a new DataPtr instances because the DataPtr class does not support set
+ // its data or context directly
+ c10::DataPtr shade_data_ptr(
+ new_shade_data_context->cpu_raw_data,
+ new_shade_data_context,
+ &(cpu::ShadeDataContext::freeShadeDataContext),
+ ipexTensor.device().type());
+
+ ipexTensor.unsafeGetTensorImpl()->storage().set_data_ptr(std::move(shade_data_ptr));
+
+ if (dstDesc.is_plain()) {
+ cpu::dbl::comm::sync_shape_from_dil_to_aten(ipexTensor, dst);
+ }
+}
 
 void attachShadeDataContext(const at::Tensor& tensor) {
  auto tensor_storage_impl = tensor.storage().unsafeGetStorageImpl();

torch_ipex/csrc/aten_ipex_bridge.h

@@ -3,6 +3,7 @@
 #include <ATen/Device.h>
 #include <ATen/Functions.h>
 #include <ATen/Tensor.h>
+#include "cpu/dbl/Common.h"
 
 #include <vector>
 
@@ -22,6 +23,13 @@ void attachShadeDataContext(const at::Tensor& tensor);
  */
 void reorderDilTensorToPublic(const at::Tensor& ipexTensor);
 
+/**
+ * Reorder to a DNNL tensor with specified descriptor no matter input tensor is a DNNL tensor or not
+ * 
+ * @param[in] ipexTensor The input tensor to be reordered to the spcified DNNL descriptor
+ */
+void reorderDilTensorGeneric(const at::Tensor& ipexTensor, const dil::tensor::desc& dstDesc);
+
 /**
  * Reorder the input tensor to the specified scalar type. It is an optimized version for
  * DNNL OP. It means that if DNNL supports current OP, you should call this API. Otherwise, you

torch_ipex/csrc/cpu/DevOPs.cpp

@@ -859,7 +859,7 @@ at::Tensor AtenIpexCPUDev::dil_adaptive_avg_pool2d(
  DEBUG("AtenIpexCPUDev::dil_adaptive_avg_pool2d\n");
  CHECK_DNNL_OP_PRE_COND(input);
  auto output_size_vec =
- dbl::pool::expand_param_if_needed(output_size, "output_size", input.dim() - 2);
+ dbl::comm::expand_param_if_needed(output_size, "output_size", input.dim() - 2);
  std::vector<int64_t> kernel_size(input.dim() - 2);
  for (int64_t i = 2; i < input.dim(); ++i) {
  auto s1 = input.size(i);

torch_ipex/csrc/cpu/FusionOPs.cpp

@@ -20,6 +20,8 @@
 namespace torch_ipex {
 namespace cpu {
 
+using namespace dbl::comm;
+
 at::Tensor AtenIpexJITDev::dil_convolution_relu(
  const at::Tensor & input,
  const at::Tensor & weight,
@@ -35,11 +37,11 @@ at::Tensor AtenIpexJITDev::dil_convolution_relu(
  auto input_contiguous = input.contiguous();
  auto weight_contiguous = weight.contiguous();
 
- dil_input = dbl::comm::try_gen_dil_tensor(input_contiguous);
- dil_weight = dbl::comm::try_gen_dil_tensor(weight_contiguous);
+ dil_input = try_gen_dil_tensor(input_contiguous);
+ dil_weight = try_gen_dil_tensor(weight_contiguous);
  if (bias.defined()) {
  auto bias_contiguous = bias.contiguous();
- dil_bias = dbl::comm::try_gen_dil_tensor(bias_contiguous);
+ dil_bias = try_gen_dil_tensor(bias_contiguous);
  }
 
  dil::tensor dil_output = dbl::conv::conv2d_impl(
@@ -52,7 +54,7 @@ at::Tensor AtenIpexJITDev::dil_convolution_relu(
  groups,
  dil::attr_t::fuse_relu());
 
- return dbl::comm::gen_aten_tensor_by(std::move(dil_output));
+ return gen_aten_tensor_by(std::move(dil_output));
 }
 
 static at::Tensor& dil_convolution_inplace_fusion(
@@ -74,12 +76,12 @@ static at::Tensor& dil_convolution_inplace_fusion(
  auto weight_contiguous = weight.contiguous();
  auto output_contiguous = accumu.contiguous();
 
- dil_input = dbl::comm::try_gen_dil_tensor(input_contiguous);
- dil_weight = dbl::comm::try_gen_dil_tensor(weight_contiguous);
- dil_output = dbl::comm::try_gen_dil_tensor(output_contiguous);
+ dil_input = try_gen_dil_tensor(input_contiguous);
+ dil_weight = try_gen_dil_tensor(weight_contiguous);
+ dil_output = try_gen_dil_tensor(output_contiguous);
  if (bias.defined()) {
  auto bias_contiguous = bias.contiguous();
- dil_bias = dbl::comm::try_gen_dil_tensor(bias_contiguous);
+ dil_bias = try_gen_dil_tensor(bias_contiguous);
  }
 
  dbl::conv::conv2d_inplace_impl(
@@ -93,7 +95,7 @@ static at::Tensor& dil_convolution_inplace_fusion(
  groups,
  attr);
 
- dbl::comm::sync_shape_from_dil_to_aten(accumu, dil_output);
+ sync_shape_from_dil_to_aten(accumu, dil_output);
  return accumu;
 }
 

torch_ipex/csrc/cpu/Prepack.cpp

@@ -0,0 +1,37 @@
+#include "Prepack.h"
+#include "dbl/Common.h"
+#include "torch_ipex/csrc/aten_ipex_bridge.h"
+#include "torch_ipex/csrc/utils.h"
+
+namespace torch_ipex {
+
+using namespace cpu::dbl::comm;
+
+void AtenIpexPrepack::prepack_conv_weight(
+ at::Tensor &weight,
+ at::IntArrayRef padding,
+ at::IntArrayRef stride,
+ at::IntArrayRef dilation,
+ int64_t groups) {
+ TORCH_CHECK(weight.device().type() == at::DeviceType::DPCPP,
+ "Cannot prepack a non-dpcpp tensor. Call t.to('dpcpp') first.");
+
+ auto kdims = weight.dim() - 2;
+ auto stride_vec = expand_param_if_needed(stride, "stride", kdims);
+ auto padding_vec = expand_param_if_needed(padding, "padding", kdims);
+ auto dilation_vec = expand_param_if_needed(dilation, "dilation", kdims);
+
+ auto packed_desc =
+ dil::convolution_forward::expected_weights_desc(
+ weight.sizes().vec(),
+ torch_ipex::get_dil_data_type(weight.scalar_type()),
+ stride_vec,
+ padding_vec,
+ padding_vec,
+ dilation_vec,
+ groups);
+
+ bridge::reorderDilTensorGeneric(weight, packed_desc);
+}
+
+} // namespace torch_ipex

torch_ipex/csrc/cpu/Prepack.h

@@ -0,0 +1,12 @@
+#pragma once
+
+#include <ATen/Tensor.h>
+
+namespace torch_ipex {
+
+class AtenIpexPrepack {
+ public:
+ static void prepack_conv_weight(at::Tensor &weight, at::IntArrayRef padding, at::IntArrayRef stride, at::IntArrayRef dilation, int64_t groups);
+};
+
+} // namespace torch_ipex

torch_ipex/csrc/cpu/dbl/Common.cpp

@@ -73,7 +73,7 @@ at::Tensor gen_aten_tensor_by(dil::tensor&& dil_tensor) {
  nullptr,
  /*resizeable=*/false);
  auto _tensor = at::detail::make_tensor<torch_ipex::IPEXTensorImpl>(storage_impl, at::DispatchKey::DPCPPTensorId);
- dbl::comm::sync_shape_from_dil_to_aten(_tensor, shade_data_context->dil_tensor.value());
+ sync_shape_from_dil_to_aten(_tensor, shade_data_context->dil_tensor.value());
  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(_tensor.layout() == c10::kStrided);
  return _tensor;
 }
@@ -101,6 +101,23 @@ void sync_shape_from_dil_to_aten(const at::Tensor& ipex_tensor, const dil::tenso
  }
 }
 
+std::vector<int64_t> expand_param_if_needed(
+ at::IntArrayRef list_param,
+ const char* param_name,
+ int64_t expected_dim) {
+ if (list_param.size() == 1) {
+ return std::vector<int64_t>(expected_dim, list_param[0]);
+ } else if ((int64_t)list_param.size() != expected_dim) {
+ std::ostringstream ss;
+ ss << "expected " << param_name << " to be a single integer value or a "
+ << "list of " << expected_dim << " values to match the convolution "
+ << "dimensions, but got " << param_name << "=" << list_param;
+ AT_ERROR(ss.str());
+ } else {
+ return list_param.vec();
+ }
+}
+
 } // namespace comm
 } // namespace dbl
 } // namespace cpu

torch_ipex/csrc/cpu/dbl/Common.h

@@ -15,6 +15,8 @@ dil::tensor try_gen_dil_tensor(const at::Tensor &input);
 at::Tensor gen_aten_tensor_by(dil::tensor&& tensor);
 at::Tensor empty_dil_tensor(at::IntArrayRef sizes, const at::TensorOptions& options);
 void sync_shape_from_dil_to_aten(const at::Tensor& ipex_tensor, const dil::tensor &dil_tensor);
+std::vector<int64_t> expand_param_if_needed(
+ at::IntArrayRef list_param, const char *param_name, int64_t expected_dim);
 
 } // namespace comm
 } // namespace dbl