Added slice BF16/FP32 FWD/BWD kernels (#34332)

jakpiase · web-flow · commit 070cab11cf02 · 2021-09-01T15:14:47.000+02:00
* aded slice FWD FP32

* added tests for slice FWD FP32

* added slice bwd

* added bf16 tests

* CI fix

* CI fix

* added reason to skip_if

* minor change

* temporary fix for failing test

* temporary fix

* changes after review

* CI rerun
diff --git a/paddle/fluid/operators/mkldnn/slice_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/slice_mkldnn_op.cc
@@ -0,0 +1,175 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/platform/mkldnn_reuse.h"
+
+namespace paddle {
+namespace operators {
+
+using paddle::framework::Tensor;
+
+template <typename T>
+class SliceMKLDNNKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& ctx) const override {
+    this->RunKernel(ctx);
+  }
+
+  void RunKernel(const framework::ExecutionContext& ctx) const {
+    const auto& dev_ctx =
+        ctx.template device_context<platform::MKLDNNDeviceContext>();
+    const auto& onednn_engine = dev_ctx.GetEngine();
+
+    auto* x = ctx.Input<Tensor>("Input");
+    auto* out = ctx.Output<Tensor>("Out");
+
+    auto x_vec_dims = framework::vectorize(x->dims());
+    auto out_vec_dims = framework::vectorize(out->dims());
+
+    auto axes_int = ctx.Attr<std::vector<int>>("axes");
+    auto starts_int = ctx.Attr<std::vector<int>>("starts");
+    auto ends_int = ctx.Attr<std::vector<int>>("ends");
+
+    std::vector<int64_t> axes(ctx.Attr<std::vector<int>>("axes").begin(),
+                              ctx.Attr<std::vector<int>>("axes").end());
+    std::vector<int64_t> starts(ctx.Attr<std::vector<int>>("starts").begin(),
+                                ctx.Attr<std::vector<int>>("starts").end());
+    std::vector<int64_t> ends(ctx.Attr<std::vector<int>>("ends").begin(),
+                              ctx.Attr<std::vector<int>>("ends").end());
+
+    auto decrease_axis = ctx.Attr<std::vector<int>>("decrease_axis");
+
+    std::vector<int64_t> offsets(x_vec_dims.size(), 0);
+    std::vector<int64_t> slice_dims(x_vec_dims);
+
+    for (size_t i = 0; i < axes.size(); ++i) {
+      starts[i] = starts[i] < 0 ? x_vec_dims[axes[i]] + starts[i] : starts[i];
+      ends[i] = ends[i] < 0 ? x_vec_dims[axes[i]] + ends[i]
+                            : std::min(ends[i], x_vec_dims[axes[i]]);
+      offsets[axes[i]] = starts[i];
+      slice_dims[axes[i]] = ends[i] - starts[i];
+    }
+
+    mkldnn::memory::data_type x_type = framework::ToMKLDNNDataType(x->type());
+    auto key = platform::CreateKey(dev_ctx, x_vec_dims, axes, starts, ends,
+                                   x->format(), x_type);
+
+    platform::ReorderMKLDNNHandler reorder_handler(
+        x_vec_dims, x->type(), x_type, dev_ctx, onednn_engine, key);
+
+    auto reorder_src_memory_p = reorder_handler.AcquireSrcMemory(
+        x->format(), platform::to_void_cast(x->data<T>()));
+    auto slice_mem_p = reorder_handler.AcquireSubmemory(slice_dims, offsets,
+                                                        reorder_src_memory_p);
+    auto reorder_dst_memory_p = reorder_handler.AcquireDstMemory(
+        out, slice_dims, 0, x->format(), ctx.GetPlace());
+
+    auto reorder_p =
+        reorder_handler.AcquireReorder(reorder_dst_memory_p, slice_mem_p);
+    auto& astream = platform::MKLDNNDeviceContext::tls().get_stream();
+    reorder_p->execute(astream, *slice_mem_p, *reorder_dst_memory_p);
+    astream.wait();
+
+    out->set_layout(framework::DataLayout::kMKLDNN);
+    out->set_format(platform::GetMKLDNNFormat(
+        reorder_dst_memory_p->get_desc().reshape(out_vec_dims)));
+  }
+};
+
+template <typename T>
+class SliceGradMKLDNNKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& ctx) const override {
+    this->RunKernel(ctx);
+  }
+
+  void RunKernel(const framework::ExecutionContext& ctx) const {
+    const auto& dev_ctx =
+        ctx.template device_context<platform::MKLDNNDeviceContext>();
+    const auto& onednn_engine = dev_ctx.GetEngine();
+
+    auto* dout = ctx.Input<Tensor>(framework::GradVarName("Out"));
+    auto* dx = ctx.Output<Tensor>(framework::GradVarName("Input"));
+
+    auto dx_vec_dims = framework::vectorize(dx->dims());
+    auto dout_vec_dims = framework::vectorize(dout->dims());
+
+    auto axes_int = ctx.Attr<std::vector<int>>("axes");
+    auto starts_int = ctx.Attr<std::vector<int>>("starts");
+    auto ends_int = ctx.Attr<std::vector<int>>("ends");
+
+    std::vector<int64_t> axes(ctx.Attr<std::vector<int>>("axes").begin(),
+                              ctx.Attr<std::vector<int>>("axes").end());
+    std::vector<int64_t> starts(ctx.Attr<std::vector<int>>("starts").begin(),
+                                ctx.Attr<std::vector<int>>("starts").end());
+    std::vector<int64_t> ends(ctx.Attr<std::vector<int>>("ends").begin(),
+                              ctx.Attr<std::vector<int>>("ends").end());
+
+    auto decrease_axis = ctx.Attr<std::vector<int>>("decrease_axis");
+
+    std::vector<int64_t> offsets(dx_vec_dims.size(), 0);
+    std::vector<int64_t> slice_dims(dx_vec_dims);
+
+    for (size_t i = 0; i < axes.size(); ++i) {
+      starts[i] = starts[i] < 0 ? dx_vec_dims[axes[i]] + starts[i] : starts[i];
+      ends[i] = ends[i] < 0 ? dx_vec_dims[axes[i]] + ends[i]
+                            : std::min(ends[i], dx_vec_dims[axes[i]]);
+      offsets[axes[i]] = starts[i];
+      slice_dims[axes[i]] = ends[i] - starts[i];
+    }
+
+    mkldnn::memory::data_type dout_type =
+        framework::ToMKLDNNDataType(dout->type());
+    mkldnn::memory::desc md(dout_vec_dims, platform::MKLDNNGetDataType<T>(),
+                            dout->format());
+    mkldnn::memory::format_tag reorder_format_tag =
+        platform::GetMKLDNNFormat(md.reshape(slice_dims));
+
+    auto key = platform::CreateKey(dev_ctx, dout_vec_dims, axes, starts, ends,
+                                   reorder_format_tag, dout_type);
+
+    platform::ReorderMKLDNNHandler reorder_handler(
+        slice_dims, dout->type(), dout_type, dev_ctx, onednn_engine, key);
+
+    auto reorder_src_memory_p = reorder_handler.AcquireSrcMemory(
+        reorder_format_tag, platform::to_void_cast(dout->data<T>()));
+    auto reorder_dst_memory_p = reorder_handler.AcquireDstMemory(
+        dx, dx_vec_dims, 0, reorder_format_tag, ctx.GetPlace());
+    memset(dx->data<T>(), 0, reorder_dst_memory_p->get_desc().get_size());
+
+    auto slice_mem_p = reorder_handler.AcquireSubmemory(slice_dims, offsets,
+                                                        reorder_dst_memory_p);
+
+    auto reorder_p =
+        reorder_handler.AcquireReorder(slice_mem_p, reorder_src_memory_p);
+    auto& astream = platform::MKLDNNDeviceContext::tls().get_stream();
+    reorder_p->execute(astream, *reorder_src_memory_p, *slice_mem_p);
+    astream.wait();
+
+    dx->set_layout(framework::DataLayout::kMKLDNN);
+    dx->set_format(reorder_format_tag);
+  }
+};
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+REGISTER_OP_KERNEL(slice, MKLDNN, paddle::platform::CPUPlace,
+                   ops::SliceMKLDNNKernel<float>,
+                   ops::SliceMKLDNNKernel<paddle::platform::bfloat16>);
+
+namespace ops = paddle::operators;
+REGISTER_OP_KERNEL(slice_grad, MKLDNN, paddle::platform::CPUPlace,
+                   ops::SliceGradMKLDNNKernel<float>,
+                   ops::SliceGradMKLDNNKernel<paddle::platform::bfloat16>);
diff --git a/paddle/fluid/operators/mkldnn/split_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/split_mkldnn_op.cc
@@ -105,7 +105,7 @@ class SplitMKLDNNKernel : public framework::OpKernel<T> {
 
     for (size_t i = 0; i < outs_number; ++i) {
       auto out_vec_dims = framework::vectorize(outs[i]->dims());
-      auto slice_mem_p = reorder_handler.AcquireSrcSubmemory(
+      auto slice_mem_p = reorder_handler.AcquireSubmemory(
           out_vec_dims, offset, reorder_src_memory_p, i);
 
       auto reorder_dst_memory_p = reorder_handler.AcquireDstMemory(
diff --git a/paddle/fluid/operators/slice_op.cc b/paddle/fluid/operators/slice_op.cc
@@ -132,6 +132,26 @@ class SliceOp : public framework::OperatorWithKernel {
       if (platform::is_cuda_pinned_place(in_tensor.place())) {
         return framework::OpKernelType(in_tensor.type(), ctx.device_context());
       }
+
+#ifdef PADDLE_WITH_MKLDNN
+      auto input_data_type =
+          framework::OperatorWithKernel::IndicateVarDataType(ctx, "Input");
+
+      if (this->CanMKLDNNBeUsed(ctx, input_data_type)) {
+        // OneDNN uses blocking format, which cannot be always supported with
+        // reorders, because if blocked dimension is not divisible by 8 or
+        // 16(depending on which blocking format is used) submemory cannot be
+        // created, so in that scenario a fallback is needed
+        auto tmp_md = dnnl::memory::desc(
+            framework::vectorize(ctx.Input<Tensor>("Input")->dims()),
+            dnnl::memory::data_type::f32, ctx.Input<Tensor>("Input")->format());
+        if (tmp_md.data.format_desc.blocking.inner_nblks == 0)
+          return framework::OpKernelType(input_data_type, ctx.GetPlace(),
+                                         framework::DataLayout::kMKLDNN,
+                                         framework::LibraryType::kMKLDNN);
+      }
+#endif
+
       return framework::OpKernelType(in_tensor.type(), in_tensor.place());
     }
     return framework::OpKernelType(
@@ -216,6 +236,14 @@ class SliceOpMaker : public framework::OpProtoAndCheckerMaker {
         .SetDefault({});
     AddAttr<std::vector<int>>("decrease_axis", "(list<int>) decrease_axis")
         .SetDefault({});
+    AddAttr<bool>("use_mkldnn",
+                  "(bool, default false) Only used in mkldnn kernel")
+        .SetDefault(false);
+    AddAttr<std::string>(
+        "mkldnn_data_type",
+        "(string, default \"float32\"). Data type of mkldnn kernel")
+        .SetDefault("float32")
+        .InEnum({"float32", "bfloat16"});
     AddComment(R"DOC(
 Slice Operator.
 
@@ -278,12 +306,32 @@ class SliceOpGrad : public framework::OperatorWithKernel {
       ctx->SetOutputDim(x_grad_name, x_dims);
     }
   }
+
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext &ctx) const override {
-    return framework::OpKernelType(OperatorWithKernel::IndicateVarDataType(
-                                       ctx, framework::GradVarName("Out")),
-                                   ctx.device_context());
+    auto input_data_type = framework::OperatorWithKernel::IndicateVarDataType(
+        ctx, framework::GradVarName("Out"));
+
+#ifdef PADDLE_WITH_MKLDNN
+    if (this->CanMKLDNNBeUsed(ctx, input_data_type)) {
+      // OneDNN uses blocking format, which cannot be always supported with
+      // reorders, because if blocked dimension is not divisible by 8 or
+      // 16(depending on which blocking format is used) submemory cannot be
+      // created, so in that scenario a fallback is needed
+      auto tmp_md = dnnl::memory::desc(
+          framework::vectorize(
+              ctx.Input<Tensor>(framework::GradVarName("Out"))->dims()),
+          dnnl::memory::data_type::f32,
+          ctx.Input<Tensor>(framework::GradVarName("Out"))->format());
+      if (tmp_md.data.format_desc.blocking.inner_nblks == 0)
+        return framework::OpKernelType(input_data_type, ctx.GetPlace(),
+                                       framework::DataLayout::kMKLDNN,
+                                       framework::LibraryType::kMKLDNN);
+    }
+#endif
+    return framework::OpKernelType(input_data_type, ctx.GetPlace());
   }
+
   framework::OpKernelType GetKernelTypeForVar(
       const std::string &var_name, const Tensor &tensor,
       const framework::OpKernelType &expected_kernel_type) const override {
diff --git a/paddle/fluid/operators/split_op.cc b/paddle/fluid/operators/split_op.cc
@@ -78,10 +78,9 @@ class SplitOp : public framework::OperatorWithKernel {
 
 #ifdef PADDLE_WITH_MKLDNN
     if (this->CanMKLDNNBeUsed(ctx, input_data_type)) {
-      // OneDNN uses blocking format, which cannot be always
-      // supported with reorders, because if blocked dimension is not divisible
-      // by
-      // 8 or 16(depending on which blocking format is used) submemory cannot be
+      // OneDNN uses blocking format, which cannot be always supported with
+      // reorders, because if blocked dimension is not divisible by 8 or
+      // 16(depending on which blocking format is used) submemory cannot be
       // created, so in that scenario a fallback is needed
       auto tmp_md = dnnl::memory::desc(
           framework::vectorize(ctx.Input<Tensor>("X")->dims()),
diff --git a/paddle/fluid/platform/mkldnn_reuse.h b/paddle/fluid/platform/mkldnn_reuse.h
@@ -1090,9 +1090,9 @@ class ReorderMKLDNNHandler : public MKLDNNHandler {
     return this->AcquireMemory(dims_, dtype_, fmt, ptr, "@user_src_mem_p");
   }
 
-  std::shared_ptr<mkldnn::memory> AcquireSrcSubmemory(
+  std::shared_ptr<mkldnn::memory> AcquireSubmemory(
       const std::vector<int64_t>& dims, const std::vector<int64_t>& offset,
-      const std::shared_ptr<mkldnn::memory>& mem_p, int submemory_number) {
+      const std::shared_ptr<mkldnn::memory>& mem_p, int submemory_number = 0) {
     std::string local_key = key_;
     local_key.append("@submem")
         .append(std::to_string(submemory_number))
diff --git a/python/paddle/fluid/tests/unittests/mkldnn/test_slice_mkldnn_op.py b/python/paddle/fluid/tests/unittests/mkldnn/test_slice_mkldnn_op.py
diff --git a/python/paddle/fluid/tests/unittests/test_static_save_load_bf16.py b/python/paddle/fluid/tests/unittests/test_static_save_load_bf16.py
