diff --git a/imperative/python/test/integration/test_dtr.py b/imperative/python/test/integration/test_dtr.py index 56b47725..c3ea68d1 100644 --- a/imperative/python/test/integration/test_dtr.py +++ b/imperative/python/test/integration/test_dtr.py @@ -91,7 +91,7 @@ class ResNet(M.Module): def run_dtr_resnet1202(): - batch_size = 7 + batch_size = 6 resnet1202 = ResNet(BasicBlock, [200, 200, 200]) opt = optim.SGD(resnet1202.parameters(), lr=0.05, momentum=0.9, weight_decay=1e-4) gm = GradManager().attach(resnet1202.parameters()) diff --git a/imperative/src/impl/dnn_op_helper.h b/imperative/src/impl/dnn_op_helper.h index 4eb2ff72..9eede36d 100644 --- a/imperative/src/impl/dnn_op_helper.h +++ b/imperative/src/impl/dnn_op_helper.h @@ -12,6 +12,7 @@ #include "megbrain/comp_node.h" #include "megbrain/comp_node_env.h" #include "megbrain/imperative/physical_tensor.h" +#include "megbrain/rdnn/management.h" using namespace megdnn; @@ -28,13 +29,12 @@ struct DnnOprCaller { CompNode cn; DeviceTensorND dev_tensor; Workspace workspace; - std::unique_ptr op; + mgb::opr::intl::UniqPtrWithCN op; - DnnOprCaller(CompNode cn) : cn(cn), op(create_operator(cn)) {} + DnnOprCaller(CompNode cn) : cn(cn), op(std::move(create_operator(cn))) {} - static std::unique_ptr create_operator(CompNode cn) { - auto&& handle = MegDNNHandle::get(CompNodeEnv::from_comp_node(cn)).handle(); - return handle->create_operator(); + static mgb::opr::intl::UniqPtrWithCN create_operator(CompNode cn) { + return mgb::opr::intl::create_megdnn_opr(cn); } megdnn::Workspace create_workspace(TensorLayout layout) { diff --git a/imperative/src/impl/ops/batch_norm.cpp b/imperative/src/impl/ops/batch_norm.cpp index dc9cdf70..9301e14b 100644 --- a/imperative/src/impl/ops/batch_norm.cpp +++ b/imperative/src/impl/ops/batch_norm.cpp @@ -171,7 +171,7 @@ SmallVector apply_on_physical_tensor( bool empty_input = src_layout.is_empty(); size_t nr_inp = inputs.size(); - DeviceTensorND ws, reserve; + DeviceTensorND reserve; size_t sz = 0, rsz = 0; TensorLayout w_layout({sz}, dtype::Byte()); @@ -186,9 +186,7 @@ SmallVector apply_on_physical_tensor( w_layout = TensorLayout({sz}, dtype::Byte()); r_layout = TensorLayout({rsz}, dtype::Byte()); } - auto wk = Blob::make(comp_node, sz); - auto ptr = wk->storage().get(); - megdnn::Workspace dnn_wk(ptr, sz); + auto dnn_wk = dnn_opr.create_workspace(w_layout); reserve = BlobManager::inst()->alloc_workspace_with_defrag(comp_node, r_layout); // alloc memory diff --git a/imperative/src/impl/ops/convolution.cpp b/imperative/src/impl/ops/convolution.cpp index e4c19431..3db9df94 100644 --- a/imperative/src/impl/ops/convolution.cpp +++ b/imperative/src/impl/ops/convolution.cpp @@ -123,8 +123,6 @@ TensorLayout do_shape_infer( std::tuple, bool> infer_output_attrs_fallible( const OpDef& def, const SmallVector& inputs) { - auto&& conv = static_cast(def); - using Param = ::megdnn::param::Convolution; SmallVector dests(1); @@ -167,34 +165,33 @@ SmallVector apply_on_physical_tensor( inp_shapes[i] = inputs[i]->layout(); } oup_shapes[0] = out_layout; - auto&& dnn_opr = opr::intl::create_megdnn_opr(cn); - dnn_opr->param().pad_h = conv.pad_h; - dnn_opr->param().pad_w = conv.pad_w; - dnn_opr->param().stride_h = conv.stride_h; - dnn_opr->param().stride_w = conv.stride_w; - dnn_opr->param().dilate_h = conv.dilate_h; - dnn_opr->param().dilate_w = conv.dilate_w; - dnn_opr->param().sparse = conv.sparse; - dnn_opr->param().compute_mode = conv.compute_mode; - dnn_opr->param().format = conv.format; + DnnOprCaller dnn_opr(cn); + dnn_opr.op->param().pad_h = conv.pad_h; + dnn_opr.op->param().pad_w = conv.pad_w; + dnn_opr.op->param().stride_h = conv.stride_h; + dnn_opr.op->param().stride_w = conv.stride_w; + dnn_opr.op->param().dilate_h = conv.dilate_h; + dnn_opr.op->param().dilate_w = conv.dilate_w; + dnn_opr.op->param().sparse = conv.sparse; + dnn_opr.op->param().compute_mode = conv.compute_mode; + dnn_opr.op->param().format = conv.format; // shape infer TensorLayout shp({0}, inputs[0]->dtype()); shp.ndim = 0; size_t sz = setup_algo( - {inp_shapes[0], inp_shapes[1], shp, shp, oup_shapes[0]}, dnn_opr.get(), 0, - false, false, cn, conv.policy(), false); + {inp_shapes[0], inp_shapes[1], shp, shp, oup_shapes[0]}, dnn_opr.op.get(), + 0, false, false, cn, conv.policy(), false); // alloc memory DeviceTensorND bias = BlobManager::inst()->alloc_workspace_with_defrag(cn, shp); - auto wk = Blob::make(cn, sz); - auto ptr = wk->storage().get(); - megdnn::Workspace dnn_wk(ptr, sz); + TensorLayout w_layout({sz}, dtype::Byte()); + auto dnn_wk = dnn_opr.create_workspace(w_layout); // exeucte - dnn_opr->exec( + dnn_opr.op->exec( inp_tensornds[0], inp_tensornds[1], bias.as_megdnn(), bias.as_megdnn(), out.as_megdnn(), nullptr, dnn_wk); return {Tensor::make(out)}; @@ -359,7 +356,6 @@ TensorLayout do_shape_infer( std::tuple, bool> infer_output_attrs_fallible( const OpDef& def, const SmallVector& inputs) { - auto&& conv = static_cast(def); using Param = ::megdnn::param::Convolution3D; SmallVector dests(1); @@ -398,24 +394,23 @@ SmallVector apply_on_physical_tensor( inp_shapes[i] = inputs[i]->layout(); } oup_shapes[0] = out_layout; - auto&& dnn_opr = opr::intl::create_megdnn_opr(cn); - dnn_opr->param() = conv.param(); + DnnOprCaller dnn_opr(cn); + dnn_opr.op->param() = conv.param(); // shape infer size_t sz = setup_algo( - {inp_shapes[0], inp_shapes[1], oup_shapes[0]}, dnn_opr.get(), 0, false, + {inp_shapes[0], inp_shapes[1], oup_shapes[0]}, dnn_opr.op.get(), 0, false, false, cn, conv.policy(), false); // alloc memory DeviceTensorND out = BlobManager::inst()->alloc_workspace_with_defrag(cn, out_layout); - auto wk = Blob::make(cn, sz); - auto ptr = wk->storage().get(); - megdnn::Workspace dnn_wk(ptr, sz); + TensorLayout w_layout({sz}, dtype::Byte()); + auto dnn_wk = dnn_opr.create_workspace(w_layout); // exeucte - dnn_opr->exec(inp_tensornds[0], inp_tensornds[1], out.as_megdnn(), dnn_wk); + dnn_opr.op->exec(inp_tensornds[0], inp_tensornds[1], out.as_megdnn(), dnn_wk); return {Tensor::make(out)}; } diff --git a/imperative/src/impl/ops/dot.cpp b/imperative/src/impl/ops/dot.cpp index 58698a71..183e338e 100644 --- a/imperative/src/impl/ops/dot.cpp +++ b/imperative/src/impl/ops/dot.cpp @@ -29,7 +29,7 @@ SmallVector apply_on_physical_tensor( using TensorND = megdnn::TensorND; SmallVector inp_tensornds; inp_tensornds.reserve(inputs.size()); - auto&& dnn_opr = opr::intl::create_megdnn_opr(comp_node); + DnnOprCaller dnn_opr(comp_node); for (unsigned i = 0; i < inputs.size(); ++i) { auto dnn_ten = inputs[i]->dnn_tensor(); inp_tensornds.push_back(dnn_ten); @@ -37,28 +37,27 @@ SmallVector apply_on_physical_tensor( TensorLayout oup_layout{inputs[0]->dtype()}; auto inp1_tensor = inputs[0]->dnn_tensor(); auto inp2_tensor = inputs[1]->dnn_tensor(); - dnn_opr->deduce_layout(inp1_tensor.layout, inp2_tensor.layout, oup_layout); + dnn_opr.op->deduce_layout(inp1_tensor.layout, inp2_tensor.layout, oup_layout); if (inputs[0]->layout().is_empty() || inputs[1]->layout().is_empty()) { - auto fill_opr = opr::intl::create_megdnn_opr(comp_node); + DnnOprCaller fill_opr(comp_node); DeviceTensorND out = BlobManager::inst()->alloc_workspace_with_defrag(comp_node, oup_layout); - fill_opr->param() = 0; - fill_opr->exec(out.as_megdnn(), {}); + fill_opr.op->param() = 0; + fill_opr.op->exec(out.as_megdnn(), {}); return {Tensor::make(out)}; } - auto wk_size = dnn_opr->get_workspace_in_bytes( + auto sz = dnn_opr.op->get_workspace_in_bytes( inp_tensornds[0].layout, inp_tensornds[1].layout, output_descs[0].layout); DeviceTensorND out_devtensor = BlobManager::inst()->alloc_workspace_with_defrag(comp_node, oup_layout); - TensorLayout wk_layout{TensorShape{wk_size}, inputs[0]->dtype()}; - DeviceTensorND workspace = - BlobManager::inst()->alloc_workspace_with_defrag(comp_node, wk_layout); - megdnn::Workspace dnn_wk(workspace.raw_ptr(), wk_size); - dnn_opr->exec( + TensorLayout w_layout({sz}, dtype::Byte()); + auto dnn_wk = dnn_opr.create_workspace(w_layout); + + dnn_opr.op->exec( inp_tensornds[0], inp_tensornds[1], out_devtensor.as_megdnn(), dnn_wk); return {Tensor::make(out_devtensor)}; diff --git a/imperative/src/impl/ops/elemwise.cpp b/imperative/src/impl/ops/elemwise.cpp index cf3b51f2..51372629 100644 --- a/imperative/src/impl/ops/elemwise.cpp +++ b/imperative/src/impl/ops/elemwise.cpp @@ -106,9 +106,8 @@ void apply_on_device_tensornd( mgb_assert( inputs.size() == trait.arity, "%s expects %u inputs; got %zu actually", trait.name, trait.arity, inputs.size()); - auto&& dnn_opr = - opr::intl::create_megdnn_opr(inputs[0].comp_node()); - opr::Elemwise::perform(op_def.mode, (*outputs)[0], inputs, dnn_opr); + DnnOprCaller dnn_opr(inputs[0].comp_node()); + opr::Elemwise::perform(op_def.mode, (*outputs)[0], inputs, dnn_opr.op); } SmallVector apply_on_physical_tensor( @@ -139,16 +138,16 @@ SmallVector apply_on_physical_tensor( if (is_empty) { return {Tensor::make(out)}; } - auto&& dnn_opr = opr::intl::create_megdnn_opr(comp_node); + DnnOprCaller dnn_opr(comp_node); - dnn_opr->param() = op_def.param(); - if (dnn_opr->param().mode == Mode::FUSE_MUL_ADD3 || - dnn_opr->param().mode == Mode::FUSE_MUL_ADD4 || + dnn_opr.op->param() = op_def.param(); + if (dnn_opr.op->param().mode == Mode::FUSE_MUL_ADD3 || + dnn_opr.op->param().mode == Mode::FUSE_MUL_ADD4 || (inp_tensornds.size() && inp_tensornds[0].layout.dtype.category() == DTypeCategory::QUANTIZED)) { - opr::Elemwise::perform_dnn(comp_node, out, inp_tensornds, dnn_opr); + opr::Elemwise::perform_dnn(comp_node, out, inp_tensornds, dnn_opr.op); } else { - dnn_opr->exec(inp_tensornds, out.as_megdnn()); + dnn_opr.op->exec(inp_tensornds, out.as_megdnn()); } return {Tensor::make(out)}; diff --git a/imperative/src/impl/ops/misc.cpp b/imperative/src/impl/ops/misc.cpp index 8addc3bc..a099dc8c 100644 --- a/imperative/src/impl/ops/misc.cpp +++ b/imperative/src/impl/ops/misc.cpp @@ -8,6 +8,7 @@ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT * ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. */ +#include "../dnn_op_helper.h" #include "../op_trait.h" #include "megbrain/imperative/ops/autogen.h" @@ -34,8 +35,7 @@ SmallVector apply_on_physical_tensor( auto dest = outputs[size]; auto cn = dest->comp_node(); - auto&& dnn_opr = opr::intl::create_megdnn_opr(cn); - size_t wk_size = 0; + DnnOprCaller dnn_opr(cn); SmallVector srcs(size); // copy an outputs to the dnn for inplace for (size_t i = 0; i < size; ++i) { @@ -44,11 +44,11 @@ SmallVector apply_on_physical_tensor( srcs[i] = outputs[i]->dev_tensor().as_megdnn(); } megdnn::CheckNonFinite::Param param({op.scale}); - dnn_opr->param() = param; - wk_size = dnn_opr->get_workspace_in_bytes(srcs, dest->layout()); - auto wk = Blob::make(cn, wk_size); - megdnn::Workspace dnn_wk(wk->storage().get(), wk_size); - dnn_opr->exec(srcs, dest->dev_tensor().as_megdnn(), dnn_wk); + dnn_opr.op->param() = param; + size_t sz = dnn_opr.op->get_workspace_in_bytes(srcs, dest->layout()); + TensorLayout w_layout({sz}, dtype::Byte()); + auto dnn_wk = dnn_opr.create_workspace(w_layout); + dnn_opr.op->exec(srcs, dest->dev_tensor().as_megdnn(), dnn_wk); return outputs; }