From bb4222876c6b753b977f1b1aacd84c9f1d3b78c6 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Thu, 5 Aug 2021 17:58:00 +0900 Subject: [PATCH 01/13] clean up typerel --- src/relay/op/nn/nn.cc | 28 +++++++++++++++++++++++++--- src/relay/op/nn/nn.h | 23 ----------------------- 2 files changed, 25 insertions(+), 26 deletions(-) diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index a96f167df2bb..af9ede279f6e 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -236,6 +236,28 @@ Expr MakeDensePack(Expr data, Expr weight, IndexExpr units, DataType out_dtype) TVM_REGISTER_GLOBAL("relay.op.nn._make.contrib_dense_pack").set_body_typed(MakeDensePack); +bool DensePackRel(const Array& types, int num_inputs, const Attrs& attrs, + const TypeReporter& reporter) { + ICHECK_EQ(types.size(), 3); + const auto* data = types[0].as(); + const auto* weight = types[1].as(); + if (data == nullptr || weight == nullptr) return false; + + const DenseAttrs* param = attrs.as(); + ICHECK(param != nullptr); + + Array oshape = data->shape; + oshape.Set((oshape.size() - 1), weight->shape[0] * weight->shape[2]); + + DataType out_dtype = param->out_dtype; + if (out_dtype.bits() == 0) { + out_dtype = data->dtype; + } + // assign output type + reporter->Assign(types[2], TensorType(oshape, out_dtype)); + return true; +} + RELAY_REGISTER_OP("nn.contrib_dense_pack") .describe(R"code(Applies a linear transformation: :math:`Y = XW^T`. @@ -249,7 +271,8 @@ RELAY_REGISTER_OP("nn.contrib_dense_pack") .add_argument("data", "nD Tensor", "Input data.") .add_argument("weight", "3D Tensor", "Packed weight matrix.") .set_support_level(10) - .add_type_rel("DensePack", DensePackRel); + + .add_type_rel("DensePack", DensePackRel); // ------------------- relay.nn.contrib_dense_pack // relay.leaky_relu @@ -307,7 +330,6 @@ bool PReluRel(const Array& types, int num_inputs, const Attrs& attrs, return true; } -template InferCorrectLayoutOutput PReluInferCorrectLayout(const Attrs& attrs, const Array& new_in_layouts, const Array& old_in_layouts, @@ -343,7 +365,7 @@ where :math:`*` is an channelwise multiplication for each sample in the batch. .add_argument("alpha", "Tensor", "Input channelwise alpha.") .set_support_level(3) .add_type_rel("PRelu", PReluRel) - .set_attr("FInferCorrectLayout", PReluInferCorrectLayout) + .set_attr("FInferCorrectLayout", PReluInferCorrectLayout) .set_attr("FTVMCompute", [](const Attrs& attrs, const Array& inputs, const Type& out_type) { const auto* param = attrs.as(); diff --git a/src/relay/op/nn/nn.h b/src/relay/op/nn/nn.h index 3dc63b31a205..6bc21473af18 100644 --- a/src/relay/op/nn/nn.h +++ b/src/relay/op/nn/nn.h @@ -116,29 +116,6 @@ bool MatmulRel(const Array& types, int num_inputs, const Attrs& attrs, return true; } -template -bool DensePackRel(const Array& types, int num_inputs, const Attrs& attrs, - const TypeReporter& reporter) { - ICHECK_EQ(types.size(), 3); - const auto* data = types[0].as(); - const auto* weight = types[1].as(); - if (data == nullptr || weight == nullptr) return false; - - const AttrType* param = attrs.as(); - ICHECK(param != nullptr); - - Array oshape = data->shape; - oshape.Set((oshape.size() - 1), weight->shape[0] * weight->shape[2]); - - DataType out_dtype = param->out_dtype; - if (out_dtype.bits() == 0) { - out_dtype = data->dtype; - } - // assign output type - reporter->Assign(types[2], TensorType(oshape, out_dtype)); - return true; -} - template bool BatchMatmulRel(const Array& types, int num_inputs, const Attrs& attrs, const TypeReporter& reporter) { From 37bbc5833c33c37c911fbea46431e2d300092328 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Thu, 5 Aug 2021 17:58:57 +0900 Subject: [PATCH 02/13] add layout transform when input is 3D --- python/tvm/topi/x86/dense_alter_op.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/python/tvm/topi/x86/dense_alter_op.py b/python/tvm/topi/x86/dense_alter_op.py index 5e15c8bf5368..0df71e98202c 100644 --- a/python/tvm/topi/x86/dense_alter_op.py +++ b/python/tvm/topi/x86/dense_alter_op.py @@ -21,6 +21,7 @@ from tvm import te from tvm import relay from tvm import autotvm +from tvm.ir import IRModule from .dense import _default_dense_pack_config from ..utils import get_const_tuple from ..nn import dense_alter_layout @@ -39,6 +40,17 @@ def _alter_dense_layout(attrs, inputs, tinfos, out_type): relay.op.get("nn.dense"), attrs, tinfos, out_type, target ) workload = autotvm.task.get_workload(outs) + + data_type = relay.transform.InferType()(IRModule.from_expr(inputs[0]))["main"].body.checked_type + assert len(data_type.shape) in [2, 3], "Input data must be a 2D or 3D tensor." + + if len(data_type.shape) == 3: + # Only supports 2D input + in_layout = "NC%dc" % data_type.shape[2] + data = relay.layout_transform(inputs[0], in_layout, "NC") + else: + data = inputs[0] + if workload: cfg = dispatch_ctx.query(target, workload) topi_impl = workload[0] @@ -63,6 +75,6 @@ def _alter_dense_layout(attrs, inputs, tinfos, out_type): ) dispatch_ctx.update(target, new_workload, cfg) weight_transform = relay.layout_transform(inputs[1], "NK", weight_layout) - return relay.nn.contrib_dense_pack(inputs[0], weight_transform, None, out_dtype) + return relay.nn.contrib_dense_pack(data, weight_transform, None, out_dtype) return None From e522aa2d683c8f5220b9f9ef3b842c07b47969c5 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Thu, 5 Aug 2021 18:17:25 +0900 Subject: [PATCH 03/13] add test --- .../python/relay/test_pass_alter_op_layout.py | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/tests/python/relay/test_pass_alter_op_layout.py b/tests/python/relay/test_pass_alter_op_layout.py index 3cee84f6a4dd..c8b5f3f8da60 100644 --- a/tests/python/relay/test_pass_alter_op_layout.py +++ b/tests/python/relay/test_pass_alter_op_layout.py @@ -1353,6 +1353,48 @@ def alter_conv2d(attrs, inputs, tinfos, out_type): assert before.body.attrs.layout == "NCHW" +def test_alter_op_dense_packed_data(): + def before(): + x = relay.var("x", shape=(1, 32, 8, 8)) + weight = relay.var("conv2d_weight", shape=(32, 32, 3, 3)) + conv = relay.nn.conv2d(x, weight, channels=32, kernel_size=(3, 3), padding=(1, 1)) + pool = relay.nn.avg_pool2d(conv, pool_size=[8, 8], padding=[0, 0, 0, 0]) + squeeze = relay.squeeze(pool, axis=[2, 3]) + dense = relay.nn.dense(squeeze, relay.var("dense_weight", shape=(16, 32))) + return relay.Function(analysis.free_vars(dense), dense) + + def expected(): + x = relay.var("x", shape=(1, 32, 8, 8)) + conv_weight = relay.var("conv2d_weight", shape=(32, 32, 3, 3)) + dense_weight = relay.var("dense_weight", shape=(16, 32)) + conv = relay.nn.contrib_conv2d_nchwc( + relay.layout_transform(x, "NCHW", "NCHW8c"), + relay.layout_transform(conv_weight, "OIHW", "OIHW8i8o"), + channels=32, + kernel_size=(3, 3), + padding=(1, 1), + data_layout="NCHW8c", + kernel_layout="OIHW8i8o", + out_layout="NCHW8c", + ) + pool = relay.nn.avg_pool2d(conv, pool_size=[8, 8], padding=[0, 0, 0, 0], layout="NCHW8c") + squeeze = relay.squeeze(pool, axis=[2, 3]) + dense = relay.nn.contrib_dense_pack( + relay.layout_transform(squeeze, "NC8c", "NC"), + relay.layout_transform(dense_weight, "NK", "NK16n"), + out_dtype="float32", + ) + return relay.Function(analysis.free_vars(dense), dense) + + with tvm.target.Target("llvm"): + with TempOpAttr( + "nn.dense", "FTVMAlterOpLayout", topi.x86.dense_alter_op._alter_dense_layout + ): + a = run_opt_pass(before(), transform.AlterOpLayout()) + b = run_opt_pass(expected(), transform.InferType()) + assert tvm.ir.structural_equal(a, b) + + if __name__ == "__main__": test_alter_op() test_alter_return_none() @@ -1377,3 +1419,4 @@ def alter_conv2d(attrs, inputs, tinfos, out_type): test_alter_op_dense() test_alter_layout_strided_slice_axes_nhwc() test_not_inplace_modify() + test_alter_op_dense_packed_data() From b6d79b5aaf3c7a03128001b0183228acfad71603 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Thu, 5 Aug 2021 18:25:43 +0900 Subject: [PATCH 04/13] update doc to clarify that only 2D input data is supported --- python/tvm/relay/op/nn/_nn.py | 6 +++--- python/tvm/relay/op/nn/nn.py | 7 +++---- src/relay/op/nn/nn.cc | 10 +++++----- 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py index 96cef8bc3588..a9ccc5aa2d24 100644 --- a/python/tvm/relay/op/nn/_nn.py +++ b/python/tvm/relay/op/nn/_nn.py @@ -1259,9 +1259,9 @@ def dense_shape_func(attrs, inputs, _): @script def _dense_pack_shape_func(data_shape, weight_shape): out = output_tensor((data_shape.shape[0],), "int64") - for i in const_range(out.shape[0] - 1): - out[i] = data_shape[i] - out[out.shape[0] - 1] = weight_shape[0] * weight_shape[2] + assert data_shape.shape[0] == 2, "Input data must be 2D" + out[0] = data_shape[0] + out[1] = weight_shape[0] * weight_shape[2] return out diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index 64b397a4d4f9..531a35c94200 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -1550,7 +1550,7 @@ def dense(data, weight, units=None, out_dtype=""): def contrib_dense_pack(data, weight, units=None, out_dtype=""): """Dense operator. - Applies a linear transformation + Applies a linear transformation with packed weight .. math:: @@ -1560,7 +1560,7 @@ def contrib_dense_pack(data, weight, units=None, out_dtype=""): ---------- data : tvm.relay.Expr The input data to the operator, - of shape `(d_1, d_2, ..., d_n, units_in)`. + of shape `(batch, units_in)`. weight : tvm.relay.Expr The transformed weight expressions, 3-D matrix, @@ -1570,8 +1570,7 @@ def contrib_dense_pack(data, weight, units=None, out_dtype=""): Number of hidden units of the dense transformation. out_dtype : str, optional - Specifies the output data type for mixed precision dense, - of shape `(d_1, d_2, ..., d_n, units)`. + Specifies the output data type for mixed precision dense. Returns ------- diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index af9ede279f6e..1384a2f13f95 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -246,8 +246,9 @@ bool DensePackRel(const Array& types, int num_inputs, const Attrs& attrs, const DenseAttrs* param = attrs.as(); ICHECK(param != nullptr); + ICHECK(data->shape.size() == 2) << "Input data must be 2D"; Array oshape = data->shape; - oshape.Set((oshape.size() - 1), weight->shape[0] * weight->shape[2]); + oshape.Set(1, weight->shape[0] * weight->shape[2]); DataType out_dtype = param->out_dtype; if (out_dtype.bits() == 0) { @@ -261,17 +262,16 @@ bool DensePackRel(const Array& types, int num_inputs, const Attrs& attrs, RELAY_REGISTER_OP("nn.contrib_dense_pack") .describe(R"code(Applies a linear transformation: :math:`Y = XW^T`. -- **data**: `(x1, x2, ..., xn, input_dim)` +- **data**: `(batch, input_dim)` - **weight**: `(units // pack_weight_tile, input_dim, pack_weight_tile)` -- **out**: `(x1, x2, ..., xn, units)`. +- **out**: `(batch, units)`. )code" TVM_ADD_FILELINE) .set_attrs_type() .set_num_inputs(2) - .add_argument("data", "nD Tensor", "Input data.") + .add_argument("data", "2D Tensor", "Input data.") .add_argument("weight", "3D Tensor", "Packed weight matrix.") .set_support_level(10) - .add_type_rel("DensePack", DensePackRel); // ------------------- relay.nn.contrib_dense_pack From 8d30a6dd35b4c95c1459e21ae01e8b263bfe925e Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 11 Aug 2021 12:16:02 +0900 Subject: [PATCH 05/13] add weight_layout attribute in dense --- include/tvm/relay/attrs/nn.h | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h index 694001f612e7..0f9e01ebb5fd 100644 --- a/include/tvm/relay/attrs/nn.h +++ b/include/tvm/relay/attrs/nn.h @@ -992,6 +992,7 @@ struct DenseAttrs : public tvm::AttrsNode { IndexExpr units; tvm::String auto_scheduler_rewritten_layout; // The layout after auto-scheduler's layout rewrite DataType out_dtype; + tvm::String weight_layout; TVM_DECLARE_ATTRS(DenseAttrs, "relay.attrs.DenseAttrs") { TVM_ATTR_FIELD(units).describe("Number of hidden units of the dense transformation."); @@ -1000,6 +1001,9 @@ struct DenseAttrs : public tvm::AttrsNode { TVM_ATTR_FIELD(out_dtype) .set_default(NullValue()) .describe("Output data type, set to explicit type under mixed precision setting"); + TVM_ATTR_FIELD(weight_layout) + .set_default("NK") + .describe("Dimension ordering of weight. Packed layouts, such as NK8n, are possible."); } }; From 682a62ac1f50f8f9aa6b933d3416dbffbdc94297 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 11 Aug 2021 12:16:23 +0900 Subject: [PATCH 06/13] remove explicit layout transform from dense_alter_op.py --- python/tvm/topi/x86/dense_alter_op.py | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/python/tvm/topi/x86/dense_alter_op.py b/python/tvm/topi/x86/dense_alter_op.py index 0df71e98202c..cb2f1929d395 100644 --- a/python/tvm/topi/x86/dense_alter_op.py +++ b/python/tvm/topi/x86/dense_alter_op.py @@ -21,7 +21,6 @@ from tvm import te from tvm import relay from tvm import autotvm -from tvm.ir import IRModule from .dense import _default_dense_pack_config from ..utils import get_const_tuple from ..nn import dense_alter_layout @@ -41,16 +40,6 @@ def _alter_dense_layout(attrs, inputs, tinfos, out_type): ) workload = autotvm.task.get_workload(outs) - data_type = relay.transform.InferType()(IRModule.from_expr(inputs[0]))["main"].body.checked_type - assert len(data_type.shape) in [2, 3], "Input data must be a 2D or 3D tensor." - - if len(data_type.shape) == 3: - # Only supports 2D input - in_layout = "NC%dc" % data_type.shape[2] - data = relay.layout_transform(inputs[0], in_layout, "NC") - else: - data = inputs[0] - if workload: cfg = dispatch_ctx.query(target, workload) topi_impl = workload[0] @@ -74,7 +63,6 @@ def _alter_dense_layout(attrs, inputs, tinfos, out_type): topi_impl, ) dispatch_ctx.update(target, new_workload, cfg) - weight_transform = relay.layout_transform(inputs[1], "NK", weight_layout) - return relay.nn.contrib_dense_pack(data, weight_transform, None, out_dtype) + return relay.nn.contrib_dense_pack(inputs[0], inputs[1], weight_layout, None, out_dtype) return None From 4bbefe1f93e090429694ad29184c37748778de89 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 11 Aug 2021 12:21:50 +0900 Subject: [PATCH 07/13] Add DensePackInferCorrectLayout to insert layout transform --- python/tvm/relay/op/nn/nn.py | 7 +++++-- src/relay/op/nn/nn.cc | 14 +++++++++++++- tests/python/relay/test_pass_alter_op_layout.py | 5 ++++- 3 files changed, 22 insertions(+), 4 deletions(-) diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index 531a35c94200..e882bcf7e271 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -1548,7 +1548,7 @@ def dense(data, weight, units=None, out_dtype=""): return _make.dense(data, weight, units, out_dtype) -def contrib_dense_pack(data, weight, units=None, out_dtype=""): +def contrib_dense_pack(data, weight, weight_layout="NK", units=None, out_dtype=""): """Dense operator. Applies a linear transformation with packed weight @@ -1566,6 +1566,9 @@ def contrib_dense_pack(data, weight, units=None, out_dtype=""): The transformed weight expressions, 3-D matrix, of shape `(units // pack_weight_tile, units_in, pack_weight_tile)`. + weight_layout: str + The layout of weight, such as "NK" or "NK8n". + units : int, optional Number of hidden units of the dense transformation. @@ -1577,7 +1580,7 @@ def contrib_dense_pack(data, weight, units=None, out_dtype=""): result : tvm.relay.Expr The computed result. """ - return _make.contrib_dense_pack(data, weight, units, out_dtype) + return _make.contrib_dense_pack(data, weight, weight_layout, units, out_dtype) def fifo_buffer(data, buffer, axis): diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index 1384a2f13f95..9c7f5d9160d4 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -226,10 +226,12 @@ RELAY_REGISTER_OP("nn.dense") // ------------------- relay.nn.contrib_dense_pack // Positional relay function to create dense_pack operator used by frontend FFI. -Expr MakeDensePack(Expr data, Expr weight, IndexExpr units, DataType out_dtype) { +Expr MakeDensePack(Expr data, Expr weight, tvm::String weight_layout, IndexExpr units, + DataType out_dtype) { auto attrs = make_object(); attrs->units = units; attrs->out_dtype = out_dtype; + attrs->weight_layout = std::move(weight_layout); static const Op& op = Op::Get("nn.contrib_dense_pack"); return Call(op, {data, weight}, Attrs(attrs), {}); } @@ -259,6 +261,15 @@ bool DensePackRel(const Array& types, int num_inputs, const Attrs& attrs, return true; } +InferCorrectLayoutOutput DensePackInferCorrectLayout(const Attrs& attrs, + const Array& new_in_layouts, + const Array& old_in_layouts, + const Array& old_in_types) { + auto params = attrs.as(); + ICHECK(params); + return InferCorrectLayoutOutput({"NC", params->weight_layout}, {"NC"}, attrs); +} + RELAY_REGISTER_OP("nn.contrib_dense_pack") .describe(R"code(Applies a linear transformation: :math:`Y = XW^T`. @@ -272,6 +283,7 @@ RELAY_REGISTER_OP("nn.contrib_dense_pack") .add_argument("data", "2D Tensor", "Input data.") .add_argument("weight", "3D Tensor", "Packed weight matrix.") .set_support_level(10) + .set_attr("FInferCorrectLayout", DensePackInferCorrectLayout) .add_type_rel("DensePack", DensePackRel); // ------------------- relay.nn.contrib_dense_pack diff --git a/tests/python/relay/test_pass_alter_op_layout.py b/tests/python/relay/test_pass_alter_op_layout.py index c8b5f3f8da60..1607b7df4a1e 100644 --- a/tests/python/relay/test_pass_alter_op_layout.py +++ b/tests/python/relay/test_pass_alter_op_layout.py @@ -1315,7 +1315,9 @@ def expected(): weight = relay.var("weight", shape=(48, 64)) target_layout = "NK16n" weight_transform = relay.layout_transform(weight, "NK", target_layout) - y = relay.nn.contrib_dense_pack(x, weight_transform, units=None, out_dtype="float32") + y = relay.nn.contrib_dense_pack( + x, weight_transform, target_layout, units=None, out_dtype="float32" + ) y = relay.Function(analysis.free_vars(y), y) return y @@ -1382,6 +1384,7 @@ def expected(): dense = relay.nn.contrib_dense_pack( relay.layout_transform(squeeze, "NC8c", "NC"), relay.layout_transform(dense_weight, "NK", "NK16n"), + "NK16n", out_dtype="float32", ) return relay.Function(analysis.free_vars(dense), dense) From d2fbd99e23b517b88bb27316424a6c71509d9c03 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 11 Aug 2021 12:49:23 +0900 Subject: [PATCH 08/13] relax type rel --- src/relay/op/nn/nn.cc | 28 ++++++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index 9c7f5d9160d4..5e75dc6a33ef 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -206,6 +206,13 @@ Expr MakeDense(Expr data, Expr weight, IndexExpr units, DataType out_dtype) { return Call(op, {data, weight}, Attrs(attrs), {}); } +InferCorrectLayoutOutput DenseInferCorrectLayout(const Attrs& attrs, + const Array& new_in_layouts, + const Array& old_in_layouts, + const Array& old_in_types) { + return InferCorrectLayoutOutput({"NC", "NK"}, {"NC"}, attrs); +} + TVM_REGISTER_GLOBAL("relay.op.nn._make.dense").set_body_typed(MakeDense); RELAY_REGISTER_OP("nn.dense") @@ -221,6 +228,7 @@ RELAY_REGISTER_OP("nn.dense") .add_argument("data", "nD Tensor", "Input data.") .add_argument("weight", "2D Tensor", "Weight matrix.") .set_support_level(1) + .set_attr("FInferCorrectLayout", DenseInferCorrectLayout) .add_type_rel("Dense", MatmulRel); // ------------------- relay.nn.dense @@ -248,9 +256,25 @@ bool DensePackRel(const Array& types, int num_inputs, const Attrs& attrs, const DenseAttrs* param = attrs.as(); ICHECK(param != nullptr); - ICHECK(data->shape.size() == 2) << "Input data must be 2D"; + // Sicne the topi impl only supports 2D data, we want to enable the following check. + // However, this function can be called when layout transform on inputs has not been + // inserted yet during AlterOpLayout. In such cases, the input data can be packed, + // i.e. data->shape.size() == 3. + // ICHECK_EQ(data->shape.size(), 2) + Array oshape = data->shape; - oshape.Set(1, weight->shape[0] * weight->shape[2]); + if (weight->shape.size() == 3) { + // The packed case, after AlterOpLayout is complete + ICHECK_EQ(data->shape.size(), 2) << "Layout transform has been applied to weight but not to " + "data. Only 2D data is supported."; + oshape.Set((oshape.size() - 1), weight->shape[0] * weight->shape[2]); + } else { + // This code path hits when contrib_dense_pack is called but before layout transform + // on inputs are inserted. + ICHECK_EQ(weight->shape.size(), 2) + << "weight shape before layout transform is applied should be 2D."; + oshape.Set((oshape.size() - 1), weight->shape[0] * weight->shape[1]); + } DataType out_dtype = param->out_dtype; if (out_dtype.bits() == 0) { From 4c4b9a7b30709c9357c91be33d02240e780328ab Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 11 Aug 2021 12:57:07 +0900 Subject: [PATCH 09/13] revert type rel relax and add check on dim --- src/relay/op/nn/nn.cc | 20 +++----------------- 1 file changed, 3 insertions(+), 17 deletions(-) diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index 5e75dc6a33ef..c7c9b3b9e293 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -256,25 +256,11 @@ bool DensePackRel(const Array& types, int num_inputs, const Attrs& attrs, const DenseAttrs* param = attrs.as(); ICHECK(param != nullptr); - // Sicne the topi impl only supports 2D data, we want to enable the following check. - // However, this function can be called when layout transform on inputs has not been - // inserted yet during AlterOpLayout. In such cases, the input data can be packed, - // i.e. data->shape.size() == 3. - // ICHECK_EQ(data->shape.size(), 2) + ICHECK_EQ(data->shape.size(), 2) << "Only 2D data is supported"; + ICHECK_EQ(weight->shape.size(), 3) << "Weight is not packed"; Array oshape = data->shape; - if (weight->shape.size() == 3) { - // The packed case, after AlterOpLayout is complete - ICHECK_EQ(data->shape.size(), 2) << "Layout transform has been applied to weight but not to " - "data. Only 2D data is supported."; - oshape.Set((oshape.size() - 1), weight->shape[0] * weight->shape[2]); - } else { - // This code path hits when contrib_dense_pack is called but before layout transform - // on inputs are inserted. - ICHECK_EQ(weight->shape.size(), 2) - << "weight shape before layout transform is applied should be 2D."; - oshape.Set((oshape.size() - 1), weight->shape[0] * weight->shape[1]); - } + oshape.Set(1, weight->shape[0] * weight->shape[2]); DataType out_dtype = param->out_dtype; if (out_dtype.bits() == 0) { From e6da45e421fbe475e30114162e5ca2dfcfff6dfc Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 11 Aug 2021 16:50:58 +0900 Subject: [PATCH 10/13] introduce DensePackAttrs to avoid breaking dense op --- include/tvm/relay/attrs/nn.h | 17 ++++++++++++++++- src/relay/op/nn/nn.cc | 8 +++++--- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h index 0f9e01ebb5fd..77cba5fa2ff1 100644 --- a/include/tvm/relay/attrs/nn.h +++ b/include/tvm/relay/attrs/nn.h @@ -992,11 +992,26 @@ struct DenseAttrs : public tvm::AttrsNode { IndexExpr units; tvm::String auto_scheduler_rewritten_layout; // The layout after auto-scheduler's layout rewrite DataType out_dtype; - tvm::String weight_layout; TVM_DECLARE_ATTRS(DenseAttrs, "relay.attrs.DenseAttrs") { TVM_ATTR_FIELD(units).describe("Number of hidden units of the dense transformation."); + // use 0 bits to indicate none. + TVM_ATTR_FIELD(out_dtype) + .set_default(NullValue()) + .describe("Output data type, set to explicit type under mixed precision setting"); + } +}; + +/*! \brief Attributes for dense_pack operator */ +struct DensePackAttrs : public tvm::AttrsNode { + IndexExpr units; + DataType out_dtype; + tvm::String weight_layout; + + TVM_DECLARE_ATTRS(DensePackAttrs, "relay.attrs.DensePackAttrs") { + TVM_ATTR_FIELD(units).describe("Number of hidden units of the dense transformation."); + // use 0 bits to indicate none. TVM_ATTR_FIELD(out_dtype) .set_default(NullValue()) diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index c7c9b3b9e293..a05e460dc680 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -233,10 +233,12 @@ RELAY_REGISTER_OP("nn.dense") // ------------------- relay.nn.dense // ------------------- relay.nn.contrib_dense_pack +TVM_REGISTER_NODE_TYPE(DensePackAttrs); + // Positional relay function to create dense_pack operator used by frontend FFI. Expr MakeDensePack(Expr data, Expr weight, tvm::String weight_layout, IndexExpr units, DataType out_dtype) { - auto attrs = make_object(); + auto attrs = make_object(); attrs->units = units; attrs->out_dtype = out_dtype; attrs->weight_layout = std::move(weight_layout); @@ -253,7 +255,7 @@ bool DensePackRel(const Array& types, int num_inputs, const Attrs& attrs, const auto* weight = types[1].as(); if (data == nullptr || weight == nullptr) return false; - const DenseAttrs* param = attrs.as(); + const DensePackAttrs* param = attrs.as(); ICHECK(param != nullptr); ICHECK_EQ(data->shape.size(), 2) << "Only 2D data is supported"; @@ -275,7 +277,7 @@ InferCorrectLayoutOutput DensePackInferCorrectLayout(const Attrs& attrs, const Array& new_in_layouts, const Array& old_in_layouts, const Array& old_in_types) { - auto params = attrs.as(); + auto params = attrs.as(); ICHECK(params); return InferCorrectLayoutOutput({"NC", params->weight_layout}, {"NC"}, attrs); } From e48ceed108561f12e6cb17774d8c6216e605c14d Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 11 Aug 2021 18:26:05 +0900 Subject: [PATCH 11/13] try fixing arm compute lib test --- tests/python/contrib/test_arm_compute_lib/test_dense.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/contrib/test_arm_compute_lib/test_dense.py b/tests/python/contrib/test_arm_compute_lib/test_dense.py index e6620a4bc1cb..7eb69cc974c0 100644 --- a/tests/python/contrib/test_arm_compute_lib/test_dense.py +++ b/tests/python/contrib/test_arm_compute_lib/test_dense.py @@ -154,7 +154,7 @@ def _get_expected_codegen(shape, weight_shape, units, dtype, has_bias=False): { "op": "const", "name": "", - "attrs": {"shape": [[[weight_shape[0]]]], "dtype": [[bias_dtype]]}, + "attrs": {"shape": [[[1, weight_shape[0]]]], "dtype": [[bias_dtype]]}, } ) From 7b86a7f2a9ca15cb6d370242385843e7969d1813 Mon Sep 17 00:00:00 2001 From: masahi Date: Thu, 12 Aug 2021 04:08:24 +0900 Subject: [PATCH 12/13] Update tests/python/contrib/test_arm_compute_lib/test_dense.py Co-authored-by: lhutton1 <35535092+lhutton1@users.noreply.github.com> --- .../contrib/test_arm_compute_lib/test_dense.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/python/contrib/test_arm_compute_lib/test_dense.py b/tests/python/contrib/test_arm_compute_lib/test_dense.py index 7eb69cc974c0..3b66ecf55350 100644 --- a/tests/python/contrib/test_arm_compute_lib/test_dense.py +++ b/tests/python/contrib/test_arm_compute_lib/test_dense.py @@ -154,7 +154,17 @@ def _get_expected_codegen(shape, weight_shape, units, dtype, has_bias=False): { "op": "const", "name": "", - "attrs": {"shape": [[[1, weight_shape[0]]]], "dtype": [[bias_dtype]]}, +if has_bias: + bias_dtype = "int32" if dtype == "uint8" else "float32" + bias_shape = [1, weight_shape[0]] if dtype == "float32" and \ + weight_shape[0] != 1 else [weight_shape[0]] + inputs.append( + { + "op": "const", + "name": "", + "attrs": {"shape": [[bias_shape]], "dtype": [[bias_dtype]]}, + } + ) } ) From b0468f55a41e17880fb5c2598c25c64f8ee06b23 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Thu, 12 Aug 2021 04:10:11 +0900 Subject: [PATCH 13/13] formatting --- .../contrib/test_arm_compute_lib/test_dense.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/tests/python/contrib/test_arm_compute_lib/test_dense.py b/tests/python/contrib/test_arm_compute_lib/test_dense.py index 3b66ecf55350..6bdff0fdb857 100644 --- a/tests/python/contrib/test_arm_compute_lib/test_dense.py +++ b/tests/python/contrib/test_arm_compute_lib/test_dense.py @@ -150,14 +150,11 @@ def _get_expected_codegen(shape, weight_shape, units, dtype, has_bias=False): if has_bias: bias_dtype = "int32" if dtype == "uint8" else "float32" - inputs.append( - { - "op": "const", - "name": "", -if has_bias: - bias_dtype = "int32" if dtype == "uint8" else "float32" - bias_shape = [1, weight_shape[0]] if dtype == "float32" and \ - weight_shape[0] != 1 else [weight_shape[0]] + bias_shape = ( + [1, weight_shape[0]] + if dtype == "float32" and weight_shape[0] != 1 + else [weight_shape[0]] + ) inputs.append( { "op": "const", @@ -165,8 +162,6 @@ def _get_expected_codegen(shape, weight_shape, units, dtype, has_bias=False): "attrs": {"shape": [[bias_shape]], "dtype": [[bias_dtype]]}, } ) - } - ) # qnn.dense params, output if dtype == "uint8":