From 08d63dfab7149504d93841d7da40d62881122ecf Mon Sep 17 00:00:00 2001 From: Shenghang Tsai Date: Tue, 14 Dec 2021 16:58:15 +0800 Subject: [PATCH] Add custom ShapeAttr in ODS (#7023) * add ShapeAttr * refine * fix doc * refine --- ci/test/1node_op_test.sh | 2 - cmake/oneflow.cmake | 3 + oneflow/ir/include/OneFlow/OneFlowBase.td | 23 +--- oneflow/ir/include/OneFlow/OneFlowOpTraits.h | 2 +- oneflow/ir/include/OneFlow/OneFlowOps.td | 6 +- .../ir/oneflow-gen-ods/oneflow-gen-ods.cpp | 2 +- .../include/OneFlow/MLIROneFlowTranslation.h | 2 +- .../lib/OneFlow/Importer.cpp | 65 +++++----- .../lib/OneFlow/MLIROneFlowTranslation.cpp | 11 +- .../test/ops/test_multi_process.py | 111 ------------------ 10 files changed, 47 insertions(+), 180 deletions(-) delete mode 100644 python/oneflow/compatible/single_client/test/ops/test_multi_process.py diff --git a/ci/test/1node_op_test.sh b/ci/test/1node_op_test.sh index 0c3057d8365..2b0214fa403 100644 --- a/ci/test/1node_op_test.sh +++ b/ci/test/1node_op_test.sh @@ -37,5 +37,3 @@ then else echo "deadlock unsolved, skipping multi-card eager" fi - -ONEFLOW_TEST_MULTI_PROCESS=1 python3 test/ops/test_multi_process.py --failfast --verbose diff --git a/cmake/oneflow.cmake b/cmake/oneflow.cmake index 832dc054dee..9e439ce92da 100644 --- a/cmake/oneflow.cmake +++ b/cmake/oneflow.cmake @@ -334,6 +334,9 @@ if(BUILD_PYTHON) target_include_directories(oneflow_internal PRIVATE ${Python_INCLUDE_DIRS} ${Python_NumPy_INCLUDE_DIRS}) target_compile_definitions(oneflow_internal PRIVATE ONEFLOW_CMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}) + if(WITH_MLIR) + add_dependencies(check-oneflow oneflow_internal) + endif(WITH_MLIR) set(gen_pip_args "") if (BUILD_CUDA) diff --git a/oneflow/ir/include/OneFlow/OneFlowBase.td b/oneflow/ir/include/OneFlow/OneFlowBase.td index 24cb67cf09d..5591d83ed31 100644 --- a/oneflow/ir/include/OneFlow/OneFlowBase.td +++ b/oneflow/ir/include/OneFlow/OneFlowBase.td @@ -3,29 +3,18 @@ include "OneFlow/OneFlowDialect.td" include "OneFlow/OneFlowInterfaces.td" - include "mlir/IR/SymbolInterfaces.td" include "mlir/Interfaces/SideEffectInterfaces.td" -def SI32ArrayAttr : TypedArrayAttrBase { - let constBuilderCall = "$_builder.getArrayAttr(llvm::to_vector<8>(llvm::map_range(values, [this](int32_t v) -> Attribute { return builder_.getSI32IntegerAttr($0); })))"; -} +def SI32ArrayAttr : TypedArrayAttrBase {} -def SI64ArrayAttr : TypedArrayAttrBase { - let constBuilderCall = "$_builder.getArrayAttr(llvm::to_vector<8>(llvm::map_range(values, [this](int64_t v) -> Attribute { return builder_.getSI64IntegerAttr($0); })))"; -} +def SI64ArrayAttr : TypedArrayAttrBase {} -def DTArrayAttr : TypedArrayAttrBase { - let constBuilderCall = "$_builder.getArrayAttr(llvm::to_vector<8>(llvm::map_range(values, [this](auto v) -> Attribute { return DataTypeAttr::get($0); })))"; -} +def ShapeAttr : TypedArrayAttrBase {} -def ShapeArrayAttr : TypedArrayAttrBase { - let constBuilderCall = "$_builder.getArrayAttr(llvm::to_vector<8>(llvm::map_range(values, [this](auto v) -> Attribute { return DenseIntElementsAttr::get($0); })))"; -} +def DTArrayAttr : TypedArrayAttrBase {} + +def ShapeArrayAttr : TypedArrayAttrBase {} def OneFlow_IsOpConfCompatible : NativeOpTrait<"IsOpConfCompatible">; def OneFlow_IsImportCompatible : NativeOpTrait<"IsImportCompatible">; diff --git a/oneflow/ir/include/OneFlow/OneFlowOpTraits.h b/oneflow/ir/include/OneFlow/OneFlowOpTraits.h index c1b674a1bcb..715bc785cc5 100644 --- a/oneflow/ir/include/OneFlow/OneFlowOpTraits.h +++ b/oneflow/ir/include/OneFlow/OneFlowOpTraits.h @@ -114,7 +114,7 @@ class TensorSource : public TraitBase { static StringRef getNdSbpAttrName() { return "nd_sbp"; } static LogicalResult verifyTrait(Operation* op) { - if (!op->hasAttrOfType(getShapeAttrName())) { + if (!op->hasAttrOfType(getShapeAttrName())) { return op->emitError("expected operation to have attribute: " + getShapeAttrName()); } if (!op->hasAttrOfType(getDataTypeAttrName())) { diff --git a/oneflow/ir/include/OneFlow/OneFlowOps.td b/oneflow/ir/include/OneFlow/OneFlowOps.td index 05de1b2aac3..16194c28d59 100644 --- a/oneflow/ir/include/OneFlow/OneFlowOps.td +++ b/oneflow/ir/include/OneFlow/OneFlowOps.td @@ -116,7 +116,7 @@ def OneFlow_VariableOp : OneFlow_ConcreteSystemOp<"variable", [OneFlow_TensorSou let input = (ins); let output = (outs AnyType:$output); let custom_attrs = (ins - AnyI64ElementsAttr:$shape, + ShapeAttr:$shape, OptionalAttr:$data_type, DefaultValuedAttr:$model_name, DefaultValuedAttr:$l1_regularization, @@ -131,7 +131,7 @@ def OneFlow_InputOp : OneFlow_ConcreteSystemOp<"input", [OneFlow_TensorSource]> let input = (ins AnyType:$input); let output = (outs AnyType:$output); let custom_attrs = (ins - OptionalAttr:$shape, + OptionalAttr:$shape, OptionalAttr:$data_type, OptionalAttr:$is_dynamic, OptionalAttr:$nd_sbp, @@ -149,7 +149,7 @@ def OneFlow_OutputOp : OneFlow_ConcreteSystemOp<"output", [OneFlow_TensorSource] let input = (ins AnyType:$input); let output = (outs AnyType:$output); let custom_attrs = (ins - OptionalAttr:$shape, + OptionalAttr:$shape, OptionalAttr:$data_type, OptionalAttr:$is_dynamic, OptionalAttr:$nd_sbp, diff --git a/oneflow/ir/oneflow-gen-ods/oneflow-gen-ods.cpp b/oneflow/ir/oneflow-gen-ods/oneflow-gen-ods.cpp index fa9c77119dd..a70d9745e8b 100644 --- a/oneflow/ir/oneflow-gen-ods/oneflow-gen-ods.cpp +++ b/oneflow/ir/oneflow-gen-ods/oneflow-gen-ods.cpp @@ -64,7 +64,7 @@ std::string GetMLIRAttrTypeName(const AttrType& attr_type) { } else if (attr_type == ::oneflow::kAtString) { return "StrAttr"; } else if (attr_type == ::oneflow::kAtShape) { - return "AnyI64ElementsAttr"; + return "ShapeAttr"; } else if (attr_type == ::oneflow::kAtDataType) { return "OneFlow_DataType"; } else if (attr_type == ::oneflow::kAtListInt32) { diff --git a/oneflow/ir/oneflow-translate/include/OneFlow/MLIROneFlowTranslation.h b/oneflow/ir/oneflow-translate/include/OneFlow/MLIROneFlowTranslation.h index 764afe967d3..8accbe8bf82 100644 --- a/oneflow/ir/oneflow-translate/include/OneFlow/MLIROneFlowTranslation.h +++ b/oneflow/ir/oneflow-translate/include/OneFlow/MLIROneFlowTranslation.h @@ -110,7 +110,7 @@ class Importer { return GetBuilder().getArrayAttr(attrs); } - DenseIntElementsAttr DenseIntElementsAttrFromShape(const ::oneflow::ShapeProto& shape); + ArrayAttr GetAttrFromShape(const ::oneflow::ShapeProto& shape); llvm::Optional GetTypeFromOneFlowDataType(::oneflow::DataType dt); OpBuilder& GetBuilder() { return builder_; } MLIRContext* GetMLIRContext() { return context_; } diff --git a/oneflow/ir/oneflow-translate/lib/OneFlow/Importer.cpp b/oneflow/ir/oneflow-translate/lib/OneFlow/Importer.cpp index e22d0208efd..fc1d35c9b50 100644 --- a/oneflow/ir/oneflow-translate/lib/OneFlow/Importer.cpp +++ b/oneflow/ir/oneflow-translate/lib/OneFlow/Importer.cpp @@ -190,17 +190,14 @@ llvm::Optional GetDataTypeAttr(MLIRContext* context } } -DenseIntElementsAttr Importer::DenseIntElementsAttrFromShape(const ::oneflow::ShapeProto& shape) { - ArrayRef values = {shape.dim().begin(), shape.dim().end()}; - RankedTensorType tt = RankedTensorType::get({static_cast(values.size())}, - GetBuilder().getIntegerType(64, true)); - ; - return DenseIntElementsAttr::get(tt, values); +ArrayAttr Importer::GetAttrFromShape(const ::oneflow::ShapeProto& shape) { + return GetBuilder().getArrayAttr(llvm::to_vector<8>(llvm::map_range( + shape.dim(), [this](int64_t v) -> Attribute { return getSI64IntegerAttr(v); }))); } -void WriteDenseIntElementsToShape(mlir::Attribute& attr, ::oneflow::ShapeProto* shape) { - for (auto int_v : attr.dyn_cast().getValues()) { - shape->add_dim(int_v); +void WriteAttrToShape(mlir::Attribute& attr, ::oneflow::ShapeProto* shape) { + for (auto v : attr.dyn_cast().getValue()) { + shape->add_dim(v.dyn_cast().getSInt()); } } @@ -235,8 +232,7 @@ LogicalResult Importer::namedAttributesFromUserOp(const ::oneflow::OperatorConf& DEFINE_ONE_ELIF(at_string, getStringAttr) #undef DEFINE_ONE_ELIF else if (value.has_at_shape()) { - attr_vec.emplace_back( - GetBuilder().getNamedAttr(name, DenseIntElementsAttrFromShape(value.at_shape()))); + attr_vec.emplace_back(GetBuilder().getNamedAttr(name, GetAttrFromShape(value.at_shape()))); } #define DEFINE_ONE_ELIF(at_key, get_attr, field) \ else if (value.has_##at_key()) { \ @@ -276,9 +272,9 @@ LogicalResult Importer::namedAttributesFromUserOp(const ::oneflow::OperatorConf& name, GetBuilder().getArrayAttr(llvm::to_vector<8>(dt_attr_list)))); } else if (value.has_at_list_shape()) { - auto dense_attr_list = llvm::map_range( - value.at_list_shape().val(), - [&](const ::oneflow::ShapeProto& s) { return DenseIntElementsAttrFromShape(s); }); + auto dense_attr_list = + llvm::map_range(value.at_list_shape().val(), + [&](const ::oneflow::ShapeProto& s) { return GetAttrFromShape(s); }); std::vector dense_attr_vector{dense_attr_list.begin(), dense_attr_list.end()}; attr_vec.emplace_back( @@ -738,7 +734,7 @@ LogicalResult Importer::ConvertUserOpAttributes(Operation* op, } else if (attr_type == ::oneflow::kAtString) { user_attr.set_at_string(attr.dyn_cast().getValue().str()); } else if (attr_type == ::oneflow::kAtShape) { - WriteDenseIntElementsToShape(attr, user_attr.mutable_at_shape()); + WriteAttrToShape(attr, user_attr.mutable_at_shape()); } else if (attr_type == ::oneflow::kAtDataType) { ::oneflow::DataType dt = ::oneflow::kInvalidDataType; if (succeeded(ConvertDT(attr, dt))) { @@ -777,11 +773,9 @@ LogicalResult Importer::ConvertUserOpAttributes(Operation* op, } } } else if (attr_type == ::oneflow::kAtListShape) { - for (auto s : attr.dyn_cast().getValue()) { + for (auto shape_attr : attr.dyn_cast().getValue()) { ::oneflow::ShapeProto* shape_ptr = user_attr.mutable_at_list_shape()->add_val(); - for (auto int_v : s.dyn_cast().getValues()) { - shape_ptr->mutable_dim()->Add(int_v); - } + WriteAttrToShape(shape_attr, shape_ptr); } } else if (attr_type == ::oneflow::kAtListString) { // attr like nd_sbp requires the existence of list even it is empty @@ -823,11 +817,12 @@ LogicalResult ConvertVariableOpConf(Operation* op, oneflow::VariableOpAdaptor& a auto* var_op_conf = op_conf->mutable_variable_conf(); var_op_conf->set_out("out"); - for (const auto& elem : adaptor.shape()) { - var_op_conf->mutable_shape()->mutable_dim()->Add(elem.getSExtValue()); + if (auto shape_attr = + op->getAttrOfType(OpTrait::TensorSource::getShapeAttrName())) { + WriteAttrToShape(shape_attr, var_op_conf->mutable_shape()); } - if (op->hasAttr("data_type")) { + if (op->hasAttr(OpTrait::TensorSource::getDataTypeAttrName())) { ::oneflow::DataType dt = ::oneflow::DataType::kInvalidDataType; if (failed(ConvertDT(adaptor.data_type(), dt))) { return failure(); } var_op_conf->set_data_type(dt); @@ -875,23 +870,22 @@ LogicalResult ConvertInputOpConf(Operation* op, oneflow::InputOpAdaptor& adaptor auto* input_op_conf = op_conf->mutable_input_conf(); input_op_conf->set_out("out"); - if (op->hasAttr("shape")) { - for (auto elem : adaptor.shape()) { - input_op_conf->mutable_blob_conf()->mutable_shape()->add_dim(elem.getSExtValue()); - } + if (auto shape_attr = + op->getAttrOfType(OpTrait::TensorSource::getShapeAttrName())) { + WriteAttrToShape(shape_attr, input_op_conf->mutable_blob_conf()->mutable_shape()); } - if (op->hasAttr("data_type")) { + if (op->hasAttr(OpTrait::TensorSource::getDataTypeAttrName())) { ::oneflow::DataType dt = ::oneflow::DataType::kInvalidDataType; if (failed(ConvertDT(adaptor.data_type(), dt))) { return failure(); } input_op_conf->mutable_blob_conf()->set_data_type(dt); } - if (op->hasAttr("is_dynamic")) { + if (op->hasAttr(OpTrait::TensorSource::getIsDynamicAttrName())) { input_op_conf->mutable_blob_conf()->set_is_dynamic(adaptor.is_dynamic().getValue()); } - if (op->hasAttr("nd_sbp")) { + if (op->hasAttr(OpTrait::TensorSource::getNdSbpAttrName())) { if (failed(ParseNdSbpFromAttr(adaptor.nd_sbp(), input_op_conf->mutable_blob_conf()->mutable_nd_sbp()))) { return failure(); @@ -919,23 +913,22 @@ LogicalResult ConvertOutputOpConf(Operation* op, oneflow::OutputOpAdaptor& adapt auto* output_op_conf = op_conf->mutable_output_conf(); output_op_conf->set_out("out"); - if (op->hasAttr("shape")) { - for (auto elem : adaptor.shape()) { - output_op_conf->mutable_blob_conf()->mutable_shape()->add_dim(elem.getSExtValue()); - } + if (auto shape_attr = + op->getAttrOfType(OpTrait::TensorSource::getShapeAttrName())) { + WriteAttrToShape(shape_attr, output_op_conf->mutable_blob_conf()->mutable_shape()); } - if (op->hasAttr("data_type")) { + if (op->hasAttr(OpTrait::TensorSource::getDataTypeAttrName())) { ::oneflow::DataType dt = ::oneflow::DataType::kInvalidDataType; if (failed(ConvertDT(adaptor.data_type(), dt))) { return failure(); } output_op_conf->mutable_blob_conf()->set_data_type(dt); } - if (op->hasAttr("is_dynamic")) { + if (op->hasAttr(OpTrait::TensorSource::getIsDynamicAttrName())) { output_op_conf->mutable_blob_conf()->set_is_dynamic(adaptor.is_dynamic().getValue()); } - if (op->hasAttr("nd_sbp")) { + if (op->hasAttr(OpTrait::TensorSource::getNdSbpAttrName())) { if (failed(ParseNdSbpFromAttr(adaptor.nd_sbp(), output_op_conf->mutable_blob_conf()->mutable_nd_sbp()))) { return failure(); diff --git a/oneflow/ir/oneflow-translate/lib/OneFlow/MLIROneFlowTranslation.cpp b/oneflow/ir/oneflow-translate/lib/OneFlow/MLIROneFlowTranslation.cpp index b0ca3e5e847..e110f493f44 100644 --- a/oneflow/ir/oneflow-translate/lib/OneFlow/MLIROneFlowTranslation.cpp +++ b/oneflow/ir/oneflow-translate/lib/OneFlow/MLIROneFlowTranslation.cpp @@ -259,8 +259,7 @@ LogicalResult JobImporter::ProcessVariableOp(const ::oneflow::OperatorConf& op_c attr_vec.emplace_back(GetBuilder().getNamedAttr( OpTrait::IsImportCompatible::getOutputLBNsAttr(), output_lbns_attr)); // attr shape - auto shape_attr = GetBuilder().getI64VectorAttr( - {op_conf.variable_conf().shape().dim().begin(), op_conf.variable_conf().shape().dim().end()}); + auto shape_attr = GetAttrFromShape(op_conf.variable_conf().shape()); auto shape_named_attr = GetBuilder().getNamedAttr(OpTrait::TensorSource::getShapeAttrName(), shape_attr); attr_vec.emplace_back(shape_named_attr); @@ -361,9 +360,7 @@ LogicalResult JobImporter::ProcessInputOp(const ::oneflow::OperatorConf& op_conf OpTrait::IsImportCompatible::getOutputLBNsAttr(), output_lbns_attr)); // attr shape if (op_conf.input_conf().blob_conf().has_shape()) { - auto shape_attr = - GetBuilder().getI64VectorAttr({op_conf.input_conf().blob_conf().shape().dim().begin(), - op_conf.input_conf().blob_conf().shape().dim().end()}); + auto shape_attr = GetAttrFromShape(op_conf.input_conf().blob_conf().shape()); attr_vec.emplace_back( GetBuilder().getNamedAttr(OpTrait::TensorSource::getShapeAttrName(), shape_attr)); } @@ -446,9 +443,7 @@ LogicalResult JobImporter::ProcessOutputOp(const ::oneflow::OperatorConf& op_con OpTrait::IsImportCompatible::getOutputLBNsAttr(), output_lbns_attr)); // attr shape if (op_conf.output_conf().blob_conf().has_shape()) { - auto shape_attr = - GetBuilder().getI64VectorAttr({op_conf.output_conf().blob_conf().shape().dim().begin(), - op_conf.output_conf().blob_conf().shape().dim().end()}); + auto shape_attr = GetAttrFromShape(op_conf.output_conf().blob_conf().shape()); attr_vec.emplace_back( GetBuilder().getNamedAttr(OpTrait::TensorSource::getShapeAttrName(), shape_attr)); } diff --git a/python/oneflow/compatible/single_client/test/ops/test_multi_process.py b/python/oneflow/compatible/single_client/test/ops/test_multi_process.py deleted file mode 100644 index d921a5e3ff0..00000000000 --- a/python/oneflow/compatible/single_client/test/ops/test_multi_process.py +++ /dev/null @@ -1,111 +0,0 @@ -""" -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 os -import unittest - -import oneflow.compatible.single_client.unittest -from oneflow.compatible import single_client as flow - - -@unittest.skipIf(flow.sysconfig.has_rpc_backend_grpc() == False, "lacks grpc") -@flow.unittest.skip_unless_1n4d() -@unittest.skipIf( - os.getenv("ONEFLOW_TEST_GITHUB_HOSTED"), - "this will fail because github hosted VM has only two CPU cores", -) -class TestMultiProcess(flow.unittest.TestCase): - def test_multi_process(test_case): - flow.config.gpu_device_num(4) - func_config = flow.FunctionConfig() - func_config.concurrency_width(1) - - @flow.global_function() - def Foo(): - with flow.scope.placement("gpu", "0:0-3"): - x = flow.get_variable( - "x", - shape=(2, 5), - dtype=flow.float, - initializer=flow.random_uniform_initializer(minval=0, maxval=1), - trainable=False, - ) - return x - - of_ret = Foo().get() - test_case.assertEqual(of_ret.numpy().shape, (2, 5)) - - def test_worker_to_master_communication(test_case): - flow.config.gpu_device_num(4) - func_config = flow.FunctionConfig() - func_config.concurrency_width(1) - - @flow.global_function() - def Foo(): - with flow.scope.placement("gpu", "0:0"): - x = flow.get_variable( - "x", - shape=(2, 5), - dtype=flow.float, - initializer=flow.random_uniform_initializer(minval=0, maxval=1), - trainable=False, - ) - with flow.scope.placement("gpu", "0:3"): - y = flow.get_variable( - "y", - shape=(2, 5), - dtype=flow.float, - initializer=flow.constant_initializer(0), - trainable=False, - ) - flow.assign(y, x) - return y - - of_ret = Foo().get() - test_case.assertEqual(of_ret.numpy().shape, (2, 5)) - - def test_worker_to_worker_communication(test_case): - flow.config.gpu_device_num(4) - func_config = flow.FunctionConfig() - func_config.concurrency_width(1) - - @flow.global_function() - def Foo(): - with flow.scope.placement("gpu", "0:1"): - x = flow.get_variable( - "x", - shape=(2, 5), - dtype=flow.float, - initializer=flow.random_uniform_initializer(minval=0, maxval=1), - trainable=False, - ) - with flow.scope.placement("gpu", "0:2"): - y = flow.get_variable( - "y", - shape=(2, 5), - dtype=flow.float, - initializer=flow.constant_initializer(0), - trainable=False, - ) - flow.assign(y, x) - return y - - of_ret = Foo().get() - test_case.assertEqual(of_ret.numpy().shape, (2, 5)) - - -if __name__ == "__main__": - unittest.main()