From 7e665c589c802904896dc9379630a350c050d8bd Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Sat, 11 Mar 2023 12:17:44 +0100 Subject: [PATCH 1/7] [TVMC] Allow passing multiple layouts to --desired-layout argument --- python/tvm/driver/tvmc/transform.py | 26 ++++++++---- tests/python/driver/tvmc/test_transform.py | 48 +++++++++++++++++++++- 2 files changed, 64 insertions(+), 10 deletions(-) diff --git a/python/tvm/driver/tvmc/transform.py b/python/tvm/driver/tvmc/transform.py index 2b34ba11b49f..aedc69586c38 100644 --- a/python/tvm/driver/tvmc/transform.py +++ b/python/tvm/driver/tvmc/transform.py @@ -103,15 +103,16 @@ def convert_to_mixed_precision(mod, ops=None, calculation_type="float16", acc_ty raise TVMCException("Error converting mixed precision : {0}".format(str(err))) -def convert_graph_layout(mod, desired_layout, ops=None): +def convert_graph_layout(mod, desired_layouts, ops=None): """Alter the layout of the input graph. Parameters ---------- mod : tvm.IRModule The relay module to convert. - desired_layout : str - The layout to convert to. + desired_layouts : list[str] + The layouts to convert to. + Expects either a single element or one str per operator. ops : list List of operators to be layout converted. @@ -123,7 +124,16 @@ def convert_graph_layout(mod, desired_layout, ops=None): if ops is None: ops = ["nn.conv2d", "nn.conv2d_transpose", "qnn.conv2d"] - desired_layouts = {op: [desired_layout, "default"] for op in ops} + assert isinstance(desired_layouts, list) and len(desired_layouts) > 0 + + if len(desired_layouts) != len(ops): + if len(desired_layouts) != 1: + raise TVMCException( + "Expected 1 or {} layouts but got {}".format(len(ops), len(desired_layouts)) + ) + desired_layouts = desired_layouts * len(ops) + + desired_layouts = {op: [desired_layouts[i], "default"] for i, op in enumerate(ops)} # Convert the layout of the graph where possible. seq = transform.Sequential( @@ -137,7 +147,7 @@ def convert_graph_layout(mod, desired_layout, ops=None): try: return seq(mod) except Exception as err: - raise TVMCException("Error converting layout to {0}: {1}".format(desired_layout, str(err))) + raise TVMCException("Error converting layouts: {}".format(str(err))) def apply_graph_transforms(mod, args): @@ -159,7 +169,7 @@ def apply_graph_transforms(mod, args): return mod # AlterLayout - if args.get("desired_layout", False): + if args.get("desired_layout", None): mod = convert_graph_layout( mod, args["desired_layout"], args.get("desired_layout_ops", None) ) @@ -211,8 +221,8 @@ def generate_transform_args(parser): parser.add_argument( "--desired-layout", choices=["NCHW", "NHWC"], - default=None, - help="Change the data layout of the whole graph.", + nargs="+", + help="Change the data layout of the graph.", ) parser.add_argument( "--desired-layout-ops", diff --git a/tests/python/driver/tvmc/test_transform.py b/tests/python/driver/tvmc/test_transform.py index 72c7cda6ff1a..bfebdccdff20 100644 --- a/tests/python/driver/tvmc/test_transform.py +++ b/tests/python/driver/tvmc/test_transform.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import pytest from unittest.mock import MagicMock import tvm @@ -23,6 +24,7 @@ from tvm.relay.expr_functor import ExprMutator from tvm.ir.instrument import pass_instrument from tvm.driver.tvmc.transform import apply_graph_transforms +from tvm.driver.tvmc.model import TVMCException def test_layout_transform_fold_constant(relay_conv2d): @@ -41,7 +43,7 @@ def run_after_pass(self, _, info): pass_names = CollectPassNames() with tvm.transform.PassContext(opt_level=3, instruments=[pass_names]): - apply_graph_transforms(relay_conv2d, {"desired_layout": desired_layout}) + apply_graph_transforms(relay_conv2d, {"desired_layout": [desired_layout]}) names = pass_names.names assert "ConvertLayout" in names @@ -61,7 +63,7 @@ def test_layout_transform_convert_layout_pass_args(relay_conv2d, monkeypatch): monkeypatch.setattr(relay.transform, "ConvertLayout", mock_convert_layout) with tvm.transform.PassContext(opt_level=3): - apply_graph_transforms(relay_conv2d, {"desired_layout": desired_layout}) + apply_graph_transforms(relay_conv2d, {"desired_layout": [desired_layout]}) mock_convert_layout.assert_called_once_with( { @@ -72,6 +74,48 @@ def test_layout_transform_convert_layout_pass_args(relay_conv2d, monkeypatch): ) +def test_layout_transform_convert_layout_pass_args_multiple(relay_conv2d, monkeypatch): + """ + Check the convert layout desired layouts arugment is what is expected when + a multiple desired layouts are provided. + """ + desired_layout = ["NHWC", "NCHW"] + desired_layout_ops = ["nn.max_pool2d", "qnn.conv2d"] + + mock_convert_layout = MagicMock() + mock_convert_layout.return_value = relay.transform.ConvertLayout({}) + monkeypatch.setattr(relay.transform, "ConvertLayout", mock_convert_layout) + + with tvm.transform.PassContext(opt_level=3): + apply_graph_transforms(relay_conv2d, {"desired_layout": desired_layout, "desired_layout_ops": desired_layout_ops}) + + mock_convert_layout.assert_called_once_with( + { + "nn.max_pool2d": ["NHWC", "default"], + "qnn.conv2d": ["NCHW", "default"], + } + ) + + +@pytest.mark.parametrize("desired", [ + (["NHWC", "NCHW"], ["nn.max_pool2d"]), + (["NHWC", "NCHW"], None), +]) +def test_layout_transform_convert_layout_pass_args_multiple_invalid(relay_conv2d, monkeypatch, desired): + """ + TODO + """ + desired_layout, desired_layout_ops = desired + + mock_convert_layout = MagicMock() + mock_convert_layout.return_value = relay.transform.ConvertLayout({}) + monkeypatch.setattr(relay.transform, "ConvertLayout", mock_convert_layout) + + with pytest.raises(TVMCException): + with tvm.transform.PassContext(opt_level=3): + apply_graph_transforms(relay_conv2d, {"desired_layout": desired_layout, "desired_layout_ops": desired_layout_ops}) + + def test_layout_transform_to_mixed_precision_pass_args_mock(relay_conv2d, monkeypatch): """ Check the mixed precision arugments which are expected when From 8d21064d73771c92901dcd45e883d3d028e8e5dc Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Sat, 11 Mar 2023 12:28:51 +0100 Subject: [PATCH 2/7] [TVMC] Allow specifying kernel layouts to --desired-layout argument --- python/tvm/driver/tvmc/transform.py | 14 +++++++++++--- tests/python/driver/tvmc/test_transform.py | 22 ++++++++++++++++++++++ 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/python/tvm/driver/tvmc/transform.py b/python/tvm/driver/tvmc/transform.py index aedc69586c38..984d041a775a 100644 --- a/python/tvm/driver/tvmc/transform.py +++ b/python/tvm/driver/tvmc/transform.py @@ -113,6 +113,7 @@ def convert_graph_layout(mod, desired_layouts, ops=None): desired_layouts : list[str] The layouts to convert to. Expects either a single element or one str per operator. + Can be only data layouts or combination of both, e.g. NHWC:HWIO ops : list List of operators to be layout converted. @@ -133,7 +134,15 @@ def convert_graph_layout(mod, desired_layouts, ops=None): ) desired_layouts = desired_layouts * len(ops) - desired_layouts = {op: [desired_layouts[i], "default"] for i, op in enumerate(ops)} + def layout_helper(layout): + if ":" in layout: + data_layout, kernel_layout = layout.split(":", 1) + else: + data_layout = layout + kernel_layout = "default" + return [data_layout, kernel_layout] + + desired_layouts = {op: layout_helper(desired_layouts[i]) for i, op in enumerate(ops)} # Convert the layout of the graph where possible. seq = transform.Sequential( @@ -220,9 +229,8 @@ def generate_transform_args(parser): # AlterLayout parser.add_argument( "--desired-layout", - choices=["NCHW", "NHWC"], nargs="+", - help="Change the data layout of the graph.", + help="Change the data/kernel layout of the graph. (i.e. NCHW or NHWC:HWIO)", ) parser.add_argument( "--desired-layout-ops", diff --git a/tests/python/driver/tvmc/test_transform.py b/tests/python/driver/tvmc/test_transform.py index bfebdccdff20..6b13a9d70d7d 100644 --- a/tests/python/driver/tvmc/test_transform.py +++ b/tests/python/driver/tvmc/test_transform.py @@ -74,6 +74,28 @@ def test_layout_transform_convert_layout_pass_args(relay_conv2d, monkeypatch): ) +def test_layout_transform_convert_kernel_layout_pass_args(relay_conv2d, monkeypatch): + """ + Check the convert layout desired layouts arugment is what is expected when + a non-default kernel layout is provided. + """ + desired_layout = "NHWC:HWIO" + desired_layout_ops = ["nn.nonv2d"] + + mock_convert_layout = MagicMock() + mock_convert_layout.return_value = relay.transform.ConvertLayout({}) + monkeypatch.setattr(relay.transform, "ConvertLayout", mock_convert_layout) + + with tvm.transform.PassContext(opt_level=3): + apply_graph_transforms(relay_conv2d, {"desired_layout": [desired_layout], "desired_layout_ops": desired_layout_ops}) + + mock_convert_layout.assert_called_once_with( + { + "nn.conv2d": ["NHWC", "HWIO"], + } + ) + + def test_layout_transform_convert_layout_pass_args_multiple(relay_conv2d, monkeypatch): """ Check the convert layout desired layouts arugment is what is expected when From eaca5d56f7d1d0b5bcd32f400960f565008febb4 Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Sat, 11 Mar 2023 20:33:19 +0100 Subject: [PATCH 3/7] lint code --- tests/python/driver/tvmc/test_transform.py | 34 ++++++++++++++++------ 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/tests/python/driver/tvmc/test_transform.py b/tests/python/driver/tvmc/test_transform.py index 6b13a9d70d7d..1a83b91ef363 100644 --- a/tests/python/driver/tvmc/test_transform.py +++ b/tests/python/driver/tvmc/test_transform.py @@ -87,7 +87,10 @@ def test_layout_transform_convert_kernel_layout_pass_args(relay_conv2d, monkeypa monkeypatch.setattr(relay.transform, "ConvertLayout", mock_convert_layout) with tvm.transform.PassContext(opt_level=3): - apply_graph_transforms(relay_conv2d, {"desired_layout": [desired_layout], "desired_layout_ops": desired_layout_ops}) + apply_graph_transforms( + relay_conv2d, + {"desired_layout": [desired_layout], "desired_layout_ops": desired_layout_ops}, + ) mock_convert_layout.assert_called_once_with( { @@ -109,7 +112,10 @@ def test_layout_transform_convert_layout_pass_args_multiple(relay_conv2d, monkey monkeypatch.setattr(relay.transform, "ConvertLayout", mock_convert_layout) with tvm.transform.PassContext(opt_level=3): - apply_graph_transforms(relay_conv2d, {"desired_layout": desired_layout, "desired_layout_ops": desired_layout_ops}) + apply_graph_transforms( + relay_conv2d, + {"desired_layout": desired_layout, "desired_layout_ops": desired_layout_ops}, + ) mock_convert_layout.assert_called_once_with( { @@ -119,13 +125,20 @@ def test_layout_transform_convert_layout_pass_args_multiple(relay_conv2d, monkey ) -@pytest.mark.parametrize("desired", [ - (["NHWC", "NCHW"], ["nn.max_pool2d"]), - (["NHWC", "NCHW"], None), -]) -def test_layout_transform_convert_layout_pass_args_multiple_invalid(relay_conv2d, monkeypatch, desired): +@pytest.mark.parametrize( + "desired", + [ + (["NHWC", "NCHW"], ["nn.max_pool2d"]), + (["NHWC", "NCHW"], None), + ], +) +def test_layout_transform_convert_layout_pass_args_multiple_invalid( + relay_conv2d, + monkeypatch, + desired, +): """ - TODO + Check invalid cases when passing multiple values to the desired layouts argument. """ desired_layout, desired_layout_ops = desired @@ -135,7 +148,10 @@ def test_layout_transform_convert_layout_pass_args_multiple_invalid(relay_conv2d with pytest.raises(TVMCException): with tvm.transform.PassContext(opt_level=3): - apply_graph_transforms(relay_conv2d, {"desired_layout": desired_layout, "desired_layout_ops": desired_layout_ops}) + apply_graph_transforms( + relay_conv2d, + {"desired_layout": desired_layout, "desired_layout_ops": desired_layout_ops}, + ) def test_layout_transform_to_mixed_precision_pass_args_mock(relay_conv2d, monkeypatch): From df17c27b846ee1cc461d1fbd1fe5cb161246b2be Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Wed, 15 Mar 2023 09:17:22 +0100 Subject: [PATCH 4/7] Address review comments --- python/tvm/driver/tvmc/transform.py | 10 +++++++++- tests/python/driver/tvmc/test_transform.py | 2 +- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/python/tvm/driver/tvmc/transform.py b/python/tvm/driver/tvmc/transform.py index 984d041a775a..14ffa8f9ae22 100644 --- a/python/tvm/driver/tvmc/transform.py +++ b/python/tvm/driver/tvmc/transform.py @@ -125,7 +125,12 @@ def convert_graph_layout(mod, desired_layouts, ops=None): if ops is None: ops = ["nn.conv2d", "nn.conv2d_transpose", "qnn.conv2d"] - assert isinstance(desired_layouts, list) and len(desired_layouts) > 0 + if not isinstance(desired_layouts, list): + # For backwards compatibility + assert isinstance(desired_layouts, str) + desired_layouts = [desired_layouts] + + assert len(desired_layouts) > 0 if len(desired_layouts) != len(ops): if len(desired_layouts) != 1: @@ -231,6 +236,9 @@ def generate_transform_args(parser): "--desired-layout", nargs="+", help="Change the data/kernel layout of the graph. (i.e. NCHW or NHWC:HWIO)", + "This option can be provided multiple times to specify per-operator layouts, " + "e.g. '--desired-layout NHWC:HWIO' (Apply same layout for every operator)." + "e.g. '--desired-layout-ops nn.conv2d nn.avg_pool2d --desired-layout NCHW NHWC'." ) parser.add_argument( "--desired-layout-ops", diff --git a/tests/python/driver/tvmc/test_transform.py b/tests/python/driver/tvmc/test_transform.py index 1a83b91ef363..06af3cb156c1 100644 --- a/tests/python/driver/tvmc/test_transform.py +++ b/tests/python/driver/tvmc/test_transform.py @@ -80,7 +80,7 @@ def test_layout_transform_convert_kernel_layout_pass_args(relay_conv2d, monkeypa a non-default kernel layout is provided. """ desired_layout = "NHWC:HWIO" - desired_layout_ops = ["nn.nonv2d"] + desired_layout_ops = ["nn.conv2d"] mock_convert_layout = MagicMock() mock_convert_layout.return_value = relay.transform.ConvertLayout({}) From 899311bd7b0bd3b2f80f425812f95709e17bb810 Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Wed, 15 Mar 2023 13:03:01 +0100 Subject: [PATCH 5/7] fix typo --- python/tvm/driver/tvmc/transform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/driver/tvmc/transform.py b/python/tvm/driver/tvmc/transform.py index 14ffa8f9ae22..bb0adaf17c7b 100644 --- a/python/tvm/driver/tvmc/transform.py +++ b/python/tvm/driver/tvmc/transform.py @@ -235,7 +235,7 @@ def generate_transform_args(parser): parser.add_argument( "--desired-layout", nargs="+", - help="Change the data/kernel layout of the graph. (i.e. NCHW or NHWC:HWIO)", + help="Change the data/kernel layout of the graph. (i.e. NCHW or NHWC:HWIO)" "This option can be provided multiple times to specify per-operator layouts, " "e.g. '--desired-layout NHWC:HWIO' (Apply same layout for every operator)." "e.g. '--desired-layout-ops nn.conv2d nn.avg_pool2d --desired-layout NCHW NHWC'." From b74b2356504925c84ee08301a5f9bed142d8cecb Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Wed, 15 Mar 2023 13:03:18 +0100 Subject: [PATCH 6/7] remove assert --- python/tvm/driver/tvmc/transform.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/tvm/driver/tvmc/transform.py b/python/tvm/driver/tvmc/transform.py index bb0adaf17c7b..be2db7208df8 100644 --- a/python/tvm/driver/tvmc/transform.py +++ b/python/tvm/driver/tvmc/transform.py @@ -130,8 +130,6 @@ def convert_graph_layout(mod, desired_layouts, ops=None): assert isinstance(desired_layouts, str) desired_layouts = [desired_layouts] - assert len(desired_layouts) > 0 - if len(desired_layouts) != len(ops): if len(desired_layouts) != 1: raise TVMCException( From d59b6ec9afa7aa35a5224dc60ce9904bb260cfee Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Wed, 15 Mar 2023 13:40:57 +0100 Subject: [PATCH 7/7] Fix lint warnings --- python/tvm/driver/tvmc/transform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/driver/tvmc/transform.py b/python/tvm/driver/tvmc/transform.py index be2db7208df8..30d9bfa639b1 100644 --- a/python/tvm/driver/tvmc/transform.py +++ b/python/tvm/driver/tvmc/transform.py @@ -236,7 +236,7 @@ def generate_transform_args(parser): help="Change the data/kernel layout of the graph. (i.e. NCHW or NHWC:HWIO)" "This option can be provided multiple times to specify per-operator layouts, " "e.g. '--desired-layout NHWC:HWIO' (Apply same layout for every operator)." - "e.g. '--desired-layout-ops nn.conv2d nn.avg_pool2d --desired-layout NCHW NHWC'." + "e.g. '--desired-layout-ops nn.conv2d nn.avg_pool2d --desired-layout NCHW NHWC'.", ) parser.add_argument( "--desired-layout-ops",