dev empty op

oneflow/python/ops/empty_op.py

@@ -0,0 +1,133 @@
+"""
+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
+from oneflow._oneflow_internal.distribute import SplitDistribute, BroadcastDistribute
+
+
+@oneflow_export("empty")
+def empty(
+    dtype: Optional[flow.dtype] = None,
+    shape: Optional[Sequence[int]] = None,
+    name: Optional[str] = None,
+    distribute: Optional[Union[SplitDistribute, BroadcastDistribute, 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.
+        distribute: The distribute attribute which can be one of SplitDistribute, BroadcastDistribute or in str format "S(N)", "B".
+
+    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
+
+        Example 1:
+
+        @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
+
+        Example 2:
+
+        flow.config.gpu_device_num(2) #set gpu num to 2
+        @flow.global_function()
+        def empty_Job() -> tp.Numpy:
+            empty_blob = flow.empty(shape=(10, 3, 3),
+                                    dtype=flow.float,
+                                    distribute="S(0)") #split at axis 0
+            return empty_blob
+
+        Example 3:
+
+        flow.config.gpu_device_num(2) #set gpu num to 2
+        @flow.global_function()
+        def empty_Job() -> tp.Numpy:
+            empty_blob = flow.empty(shape=(10, 3, 3),
+                                    dtype=flow.float,
+                                    distribute=flow.distribute.split(0)) # same as "S(0)",
+            return empty_blob
+
+    """
+
+    def _distribute_to_str(dist):
+        dist_str = ""
+        if dist is None:
+            pass
+        elif type(dist) is oneflow._oneflow_internal.distribute.SplitDistribute:
+            dist_str = "S({})".format(dist.axis)
+        elif type(dist) is oneflow._oneflow_internal.distribute.BroadcastDistribute:
+            dist_str = "B"
+        else:
+            raise ValueError("unsupported distribute")
+        return dist_str
+
+    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 = []
+    if distribute is None:
+        sbp_parallel = ""
+    elif isinstance(distribute, str):
+        sbp_parallel = distribute
+    elif isinstance(distribute, BroadcastDistribute) or isinstance(
+        distribute, SplitDistribute
+    ):
+        sbp_parallel = _distribute_to_str(distribute)
+    else:
+        raise ValueError("Wrong distribute value")
+
+    return (
+        flow.user_op_builder(name)
+        .Op("empty")
+        .Output("out")
+        .Attr("dtype", dtype)
+        .Attr("shape", shape)
+        .Attr("sbp_parallel", sbp_parallel)
+        .Build()
+        .InferAndTryRun()
+        .RemoteBlobList()[0]
+    )

oneflow/python/test/ops/test_empty.py

@@ -0,0 +1,157 @@
+"""
+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
+from oneflow._oneflow_internal.distribute import SplitDistribute, BroadcastDistribute
+
+func_config = flow.FunctionConfig()
+func_config.default_data_type(flow.float)
+
+
+def _test_empty_fp16(test_case, device_type, shape, 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, "0:0-%d" % (device_count - 1)):
+            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_empty(
+    test_case, device_type, type_name, shape, device_count, sbp_parallel=None
+):
+    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, "0:0-%d" % (device_count - 1)):
+            return flow.empty(dtype=flow_type, shape=shape, distribute=sbp_parallel)
+
+    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["device_count"] = [1]
+        for arg in GenArgList(arg_dict):
+            _test_empty(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["device_count"] = [1]
+        for arg in GenArgList(arg_dict):
+            _test_empty_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["device_count"] = [2]
+        for arg in GenArgList(arg_dict):
+            _test_empty_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["device_count"] = [2]
+        arg_dict["sbp_parallel"] = [flow.distribute.broadcast()]
+        for arg in GenArgList(arg_dict):
+            _test_empty(test_case, *arg)
+
+    def test_empty_sbp(test_case):
+        arg_dict = OrderedDict()
+        arg_dict["device_type"] = ["gpu"]
+        arg_dict["type_name_value"] = [
+            "float32",
+        ]
+        arg_dict["shape"] = [(10, 3, 4, 5)]
+        arg_dict["device_count"] = [2]
+        arg_dict["sbp_parallel"] = [
+            flow.distribute.split(0),
+            flow.distribute.broadcast(),
+            "S(0)",
+            "B",
+        ]
+        for arg in GenArgList(arg_dict):
+            _test_empty(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