dev empty op

oneflow/python/ops/empty_op.py

@@ -0,0 +1,84 @@
+"""
+Copyright 2020 The OneFlow 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.
+"""
+from __future__ import absolute_import
+
+import os
+from typing import Optional, Sequence, Union
+
+import oneflow as flow
+import oneflow.core.operator.op_conf_pb2 as op_conf_util
+import oneflow.core.register.logical_blob_id_pb2 as logical_blob_id_util
+import oneflow.python.framework.interpret_util as interpret_util
+import oneflow.python.framework.id_util as id_util
+import oneflow.python.framework.remote_blob as remote_blob_util
+from oneflow.python.oneflow_export import oneflow_export
+import oneflow._oneflow_internal
+
+
+@oneflow_export("empty")
+def empty(
+    dtype: Optional[flow.dtype] = None,
+    shape: Optional[Sequence[int]] = None,
+    name: Optional[str] = None,
+) -> oneflow._oneflow_internal.BlobDesc:
+    """This operator creates a uninitialized Blob with specified shape.
+
+    Args:
+        dtype (Optional[flow.dtype], optional): The data type of Blob. Defaults to None.
+        shape (Optional[Sequence[int]], optional): The shape of Blob. Defaults to None.
+        name (Optional[str], optional): The name for the operation. Defaults to None.
+
+    Returns:
+        oneflow._oneflow_internal.BlobDesc: The result blob.
+
+    For example:
+
+    .. code-block:: python
+
+        import oneflow as flow
+        import numpy as np
+        import oneflow.typing as tp
+
+
+        @flow.global_function()
+        def empty_Job() -> tp.Numpy:
+            empty_blob = flow.empty(shape=(1, 3, 3),
+                                    dtype=flow.float)
+            return empty_blob
+
+
+        out = empty_Job() # out tensor with shape (1, 3, 3) and data uninitialized
+
+    """
+    if name is None:
+        name = id_util.UniqueStr("Empty_")
+
+    assert dtype is not None
+
+    if shape is not None:
+        assert isinstance(shape, (list, tuple))
+    else:
+        shape = []
+    return (
+        flow.user_op_builder(name)
+        .Op("empty")
+        .Output("out")
+        .Attr("dtype", dtype)
+        .Attr("shape", shape)
+        .Build()
+        .InferAndTryRun()
+        .RemoteBlobList()[0]
+    )

oneflow/python/test/ops/test_empty.py

@@ -0,0 +1,140 @@
+"""
+Copyright 2020 The OneFlow 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.
+"""
+import unittest
+import os
+from collections import OrderedDict
+
+import numpy as np
+import oneflow as flow
+from test_util import GenArgList
+from test_util import type_name_to_flow_type
+from test_util import type_name_to_np_type
+import oneflow.typing as tp
+
+func_config = flow.FunctionConfig()
+func_config.default_data_type(flow.float)
+
+
+def _test_fp16(test_case, device_type, shape, machine_ids, device_count):
+    assert device_type in ["gpu"]
+    flow.clear_default_session()
+
+    if device_type == "cpu":
+        flow.config.cpu_device_num(device_count)
+    elif device_type == "gpu":
+        flow.config.gpu_device_num(device_count)
+
+    @flow.global_function(function_config=func_config)
+    def empty_fp16_job() -> tp.Numpy:
+        with flow.scope.placement(device_type, machine_ids):
+            out16 = flow.empty(dtype=flow.float16, shape=shape)
+            test_case.assertTrue(out16.dtype == flow.float16)
+            out32 = flow.cast(out16, flow.float32)
+            return out32
+
+    if shape == ():
+        np_shape = (1,)
+    else:
+        np_shape = shape
+    of_out = empty_fp16_job()
+    test_case.assertTrue(of_out.shape == np_shape)
+    test_case.assertTrue(of_out.dtype == np.float32)
+
+
+def _test(test_case, device_type, type_name, shape, machine_ids, device_count):
+    assert device_type in ["gpu", "cpu"]
+    flow.clear_default_session()
+
+    if device_type == "cpu":
+        flow.config.cpu_device_num(device_count)
+    elif device_type == "gpu":
+        flow.config.gpu_device_num(device_count)
+
+    flow_type = type_name_to_flow_type[type_name]
+    np_type = type_name_to_np_type[type_name]
+
+    @flow.global_function(function_config=func_config)
+    def empty_job() -> tp.Numpy:
+        with flow.scope.placement(device_type, machine_ids):
+            return flow.empty(dtype=flow_type, shape=shape)
+
+    if shape == ():
+        np_shape = (1,)
+    else:
+        np_shape = shape
+    of_out = empty_job()
+    test_case.assertTrue(of_out.shape == np_shape)
+    test_case.assertTrue(of_out.dtype == np_type)
+
+
+@flow.unittest.skip_unless_1n1d()
+class TestEmpty1n1d(flow.unittest.TestCase):
+    def test_empty(test_case):
+        arg_dict = OrderedDict()
+        arg_dict["device_type"] = ["gpu", "cpu"]
+        arg_dict["type_name_value"] = [
+            "float32",
+            "double",
+            "int8",
+            "int32",
+            "int64",
+        ]
+        arg_dict["shape"] = [(2, 3, 4, 5), (2, 3), (100, 100), (512, 512), ()]
+        arg_dict["machine_ids"] = ["0:0"]
+        arg_dict["device_count"] = [1]
+        for arg in GenArgList(arg_dict):
+            _test(test_case, *arg)
+
+    def test_empty_fp16(test_case):
+        arg_dict = OrderedDict()
+        arg_dict["device_type"] = ["gpu"]
+        arg_dict["shape"] = [(2, 3, 4, 5), (2, 3), (100, 100), (512, 512), ()]
+        arg_dict["machine_ids"] = ["0:0-0"]
+        arg_dict["device_count"] = [1]
+        for arg in GenArgList(arg_dict):
+            _test_fp16(test_case, *arg)
+
+
+@flow.unittest.skip_unless_1n2d()
+class TestEmpty1n2d(flow.unittest.TestCase):
+    def test_empty_fp16(test_case):
+        arg_dict = OrderedDict()
+        arg_dict["device_type"] = ["gpu"]
+        arg_dict["shape"] = [(2, 3, 4, 5), (2, 3), (100, 100), (512, 512), ()]
+        arg_dict["machine_ids"] = ["0:0-1"]
+        arg_dict["device_count"] = [2]
+        for arg in GenArgList(arg_dict):
+            _test_fp16(test_case, *arg)
+
+    def test_empty(test_case):
+        arg_dict = OrderedDict()
+        arg_dict["device_type"] = ["gpu", "cpu"]
+        arg_dict["type_name_value"] = [
+            "float32",
+            "double",
+            "int8",
+            "int32",
+            "int64",
+        ]
+        arg_dict["shape"] = [(2, 3, 4, 5), (2, 3), (100, 100), (512, 512), ()]
+        arg_dict["machine_ids"] = ["0:0-1"]
+        arg_dict["device_count"] = [2]
+        for arg in GenArgList(arg_dict):
+            _test(test_case, *arg)
+
+
+if __name__ == "__main__":
+    unittest.main()

oneflow/user/kernels/empty_kernel.cpp

@@ -0,0 +1,54 @@
+/*
+Copyright 2020 The OneFlow 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 "oneflow/core/framework/framework.h"
+#include "oneflow/core/kernel/new_kernel_util.h"
+
+namespace oneflow {
+namespace user_op {
+
+template<DeviceType device_type, typename T>
+class EmptyKernel final : public OpKernel {
+ public:
+  EmptyKernel() = default;
+  ~EmptyKernel() = default;
+
+ private:
+  void Compute(user_op::KernelComputeContext* ctx) const override {
+    Tensor* out_tensor = ctx->Tensor4ArgNameAndIndex("out", 0);
+    const int64_t elem_cnt = out_tensor->shape().elem_cnt();
+    CHECK_GT(elem_cnt, 0);
+
+    // Do nothing
+  }
+  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
+};
+
+#define REGISTER_EMPTY_XPU_KERNEL(device, dtype)                                           \
+  REGISTER_USER_KERNEL("empty").SetCreateFn<EmptyKernel<device, dtype>>().SetIsMatchedHob( \
+      (user_op::HobDeviceTag() == device)                                                  \
+      & (user_op::HobAttr<DataType>("dtype") == GetDataType<dtype>::value));
+
+#define REGISTER_EMPTY_KERNEL(device, dtype_pair) \
+  REGISTER_EMPTY_XPU_KERNEL(device, OF_PP_PAIR_FIRST(dtype_pair))
+
+OF_PP_SEQ_PRODUCT_FOR_EACH_TUPLE(REGISTER_EMPTY_KERNEL, DEVICE_TYPE_SEQ, ARITHMETIC_DATA_TYPE_SEQ)
+
+#ifdef WITH_CUDA
+REGISTER_EMPTY_XPU_KERNEL(DeviceType::kGPU, float16);
+#endif  // WITH_CUDA
+
+}  // namespace user_op
+}  // namespace oneflow

oneflow/user/ops/empty_op.cpp

@@ -0,0 +1,53 @@
+/*
+Copyright 2020 The OneFlow 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 "oneflow/core/framework/framework.h"
+
+namespace oneflow {
+REGISTER_USER_OP("empty")
+    .Output("out")
+    .SetOutputBufferNum(1)
+    .Attr<DataType>("dtype")
+    .Attr<Shape>("shape")
+    .SetTensorDescInferFn([](user_op::InferContext* ctx) -> Maybe<void> {
+      Shape* out_shape = ctx->Shape4ArgNameAndIndex("out", 0);
+      const Shape& shape = ctx->Attr<Shape>("shape");
+      DimVector dim_vec;
+      if (shape.NumAxes() > 0) {
+        dim_vec.insert(dim_vec.end(), shape.dim_vec().cbegin(), shape.dim_vec().cend());
+      }
+      if (dim_vec.empty()) { dim_vec.push_back(1); }
+      *out_shape = Shape(dim_vec);
+      return Maybe<void>::Ok();
+    })
+    .SetGetSbpFn([](user_op::SbpContext* ctx) -> Maybe<void> {
+      // const Shape& shape = ctx->Attr<Shape>("shape");
+      // if (shape.NumAxes() > 0) {
+      //   FOR_RANGE(int64_t, i, 0, shape.NumAxes()) {
+      //     ctx->NewBuilder().Split(ctx->outputs(), i).Build();
+      //   }
+      // }
+      ctx->NewBuilder().Broadcast(ctx->outputs()).Build();
+      ctx->NewBuilder().PartialSum(ctx->outputs()).Build();
+
+      return Maybe<void>::Ok();
+    })
+    .SetInferDataTypeFn([](user_op::InferContext* ctx) -> Maybe<void> {
+      const DataType dtype = ctx->Attr<DataType>("dtype");
+      *ctx->Dtype4ArgNameAndIndex("out", 0) = dtype;
+      return Maybe<void>::Ok();
+    });
+
+}  // namespace oneflow