record gelu data flow for int8 fusion path (#16)

Author: XiaobingSuper

intel_pytorch_extension_py/conf.py

@@ -29,7 +29,7 @@ def save(self, configure_file, default_recipe=True):
             json.dump(configures, fp, indent = 4)
 
     def get_default_recipe(self, configures):
-        elt_wise = ['relu', 'sigmoid']
+        elt_wise = ['relu', 'sigmoid', 'gelu']
         inplace_ops = ['relu_', 'add_']
         # get default recipe,
         # q+dq+conv+q+dq+relu => q+dq+conv+relu

torch_ipex/csrc/autocast_kernel.cpp

@@ -215,5 +215,15 @@ at::Tensor dropout(const at::Tensor& input, double p, bool train) {
   return at::dropout(input, p, train);
 }
 
+at::Tensor gelu(const at::Tensor& input) {
+  c10::impl::ExcludeDispatchKeyGuard no_autocastCPU(DispatchKey::AutocastCPU);
+  auto target_type = get_autocast_dtype();
+  if (at::ScalarType::Char == target_type) {
+    return int8::gelu(input);
+  }
+  // convert to fp32 path.
+  return at::gelu(cpu_cached_cast(at::kFloat, input));
+}
+
 } // autocast
 } // torch_ipex

torch_ipex/csrc/autocast_kernel.hpp

@@ -47,5 +47,7 @@ at::Tensor add_tensor(const at::Tensor& input, const at::Tensor& other, const at
 
 at::Tensor dropout(const at::Tensor& input, double p, bool train);
 
+at::Tensor gelu(const at::Tensor& input);
+
 } // autocast
 } // torch_ipex

torch_ipex/csrc/autocast_mode.cpp

@@ -181,7 +181,6 @@ MAKE_REGISTER_FUNC(ADD_NS(addbmm), "addbmm", Tensor (const Tensor &, const Tenso
 MAKE_REGISTER_FUNC(ADD_NS(convolution), "convolution", Tensor (const Tensor &, const Tensor &, const c10::optional<Tensor>&, IntArrayRef, IntArrayRef, IntArrayRef, bool, IntArrayRef, int64_t), fp32)
 MAKE_REGISTER_FUNC(ADD_NS(avg_pool2d), "avg_pool2d", Tensor (const Tensor &, IntArrayRef, IntArrayRef, IntArrayRef, bool, bool, c10::optional<int64_t>), fp32)
 MAKE_REGISTER_FUNC(ADD_NS(avg_pool3d), "avg_pool3d", Tensor (const Tensor &, IntArrayRef, IntArrayRef, IntArrayRef, bool, bool, c10::optional<int64_t>), fp32)
-MAKE_REGISTER_FUNC(ADD_NS(gelu), "gelu", Tensor (const Tensor &), fp32)
 MAKE_REGISTER_FUNC(ADD_NS(upsample_nearest1d), "upsample_nearest1d", Tensor (const Tensor &, IntArrayRef, c10::optional<double>), fp32)
 MAKE_REGISTER_FUNC(ADD_NS(upsample_nearest1d), "upsample_nearest1d.vec", Tensor (const Tensor &, c10::optional<IntArrayRef>, c10::optional<ArrayRef<double>>), fp32)
 MAKE_REGISTER_FUNC(ADD_NS(upsample_nearest2d), "upsample_nearest2d", Tensor (const Tensor &, IntArrayRef, c10::optional<double>, c10::optional<double>), fp32)
@@ -250,7 +249,7 @@ TORCH_LIBRARY_IMPL(aten, AutocastCPU, m){
   m.impl(TORCH_SELECTIVE_NAME("aten::add_.Tensor"), TORCH_FN((&torch_ipex::autocast::add_tensor_)));
   m.impl(TORCH_SELECTIVE_NAME("aten::add.Tensor"), TORCH_FN((&torch_ipex::autocast::add_tensor)));
   m.impl(TORCH_SELECTIVE_NAME("aten::dropout"), TORCH_FN((&torch_ipex::autocast::dropout)));
-
+  m.impl(TORCH_SELECTIVE_NAME("aten::gelu"), TORCH_FN((&torch_ipex::autocast::gelu)));
 }
 
 }  // namespace autocast

torch_ipex/csrc/quantization/AutoCast.cpp

@@ -637,6 +637,29 @@ at::Tensor dropout(const at::Tensor &input, double p, bool train) {
   return at::dropout(input, p, train);
 }
 
+at::Tensor gelu(const at::Tensor &input) {
+  auto op_id = torch_ipex::Int8OptConfig::fetch_and_add_ops_id();
+  if (check_int8_calibration()) {
+    auto it = tensors_flow.find(input.unsafeGetTensorImpl());
+    std::vector<std::string> op_inputs, op_outputs;
+    if (it == tensors_flow.end()) {
+      std::string op_input = "gelu." + std::to_string(op_id) + ".input";
+      op_inputs.push_back(op_input);
+    } else {
+      op_inputs.push_back(std::get<1>(it->second));
+    }
+
+    auto output = at::gelu(input);
+    std::string op_output = "gelu." + std::to_string(op_id) + ".output";
+    op_outputs.push_back(op_output);
+    tensors_flow.emplace(output.unsafeGetTensorImpl(),
+                         val_name{weakref_scales(output.getIntrusivePtr()), op_output});
+    torch_ipex::insert_or_updata_observer({input}, {output}, "gelu",
+                                          op_id, op_inputs, op_outputs);
+    return output;
+  }
+}
+  
 } // namespace autocast
 } // namespace cpu
 } // namespace torch_ipex

torch_ipex/csrc/quantization/AutoCast.hpp

@@ -50,6 +50,8 @@ at::Tensor add_tensor(const at::Tensor& input, const at::Tensor& other, const at
 
 at::Tensor dropout(const at::Tensor &input, double p, bool train);
 
+at::Tensor gelu(const at::Tensor& input);
+
 } // namespace int8
 } // namespace autocast
 } // namespace torch_ipex