From 4fc75dfc9f8bd2b37c6aafb5e92ae373e8f220de Mon Sep 17 00:00:00 2001 From: tangzhiyi Date: Thu, 25 Jan 2024 09:26:44 +0000 Subject: [PATCH 1/7] add scalar_tensor for ascendgraph --- dicp/dicp/vendor/AscendGraph/conversion.py | 4 ++++ dicp/test/ascend_scripts/ops/static.ini | 1 + dicp/test/op/test_scalar_tensor.py | 14 ++++++++++---- 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/dicp/dicp/vendor/AscendGraph/conversion.py b/dicp/dicp/vendor/AscendGraph/conversion.py index 8a58f1e19..76d9aa3b4 100644 --- a/dicp/dicp/vendor/AscendGraph/conversion.py +++ b/dicp/dicp/vendor/AscendGraph/conversion.py @@ -1483,3 +1483,7 @@ def SliceScatter(self, operand, src, dim=0, start=None, end=None, step=1): indices_expanded = indices return self.get_proxy(ascend_op.ScatterElements, (operand, indices_expanded, src, dim)) + + @register_conversion(torch.ops.aten.scalar_tensor.default) + def scalar_tensor(self, x, dtype=None, layout=None, device=None, pin_memory=None): + return self.get_const_proxy(x, dtype) diff --git a/dicp/test/ascend_scripts/ops/static.ini b/dicp/test/ascend_scripts/ops/static.ini index 688562aa5..efbd2e7cb 100644 --- a/dicp/test/ascend_scripts/ops/static.ini +++ b/dicp/test/ascend_scripts/ops/static.ini @@ -54,6 +54,7 @@ python_files = test_relu.py test_repeat.py test_rsqrt.py + test_scalar_tensor.py test_scatter.py test_select.py test_sigmoid.py diff --git a/dicp/test/op/test_scalar_tensor.py b/dicp/test/op/test_scalar_tensor.py index 19668c5bd..4b555e1f1 100644 --- a/dicp/test/op/test_scalar_tensor.py +++ b/dicp/test/op/test_scalar_tensor.py @@ -9,9 +9,14 @@ class OpModule(torch.nn.Module): - def forward(self, a, dtype): + def forward(self, a, redundant_input, dtype): + # If there is only one operator called scalar_tensor, + # ascend graph compiler will give an error: + # GE.. [Check][Param] SetInputs failed: input operator size can not be 0. + # To solve this problem, an additional redundant input is added, + # and the result of an addition operator is returned. res_default = torch.ops.aten.scalar_tensor.default(a, dtype=dtype) - return res_default + return res_default + redundant_input, res_default model = OpModule() @@ -24,9 +29,10 @@ class TestScalarTensor(): @pytest.mark.parametrize("inputs", [1.0, 3.0, 0.0]) @pytest.mark.parametrize("compiled_model", compiled_model) def test_torch_scalar_tensor(self, inputs, dtype, compiled_model): - output = model(inputs, dtype) + redundant_input = torch.randn(1, dtype=dtype) + _, output = model(inputs, redundant_input, dtype) dynamo.reset() update_dynamo_config(compiled_model.dynamic) - dicp_output = compiled_model.model(inputs, dtype) + _, dicp_output = compiled_model.model(inputs, redundant_input, dtype) assert torch.allclose(output, dicp_output.cpu(), equal_nan=True) From d490fd8e89486ae9e47b9ffcfa30a549b33ffab9 Mon Sep 17 00:00:00 2001 From: tangzhiyi Date: Thu, 25 Jan 2024 10:09:41 +0000 Subject: [PATCH 2/7] update --- dicp/dicp/dynamo_bridge/graph.py | 4 +++- dicp/test/op/test_scalar_tensor.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/dicp/dicp/dynamo_bridge/graph.py b/dicp/dicp/dynamo_bridge/graph.py index f5f5095d2..fd338c7c3 100644 --- a/dicp/dicp/dynamo_bridge/graph.py +++ b/dicp/dicp/dynamo_bridge/graph.py @@ -61,8 +61,10 @@ def make_tensor_meta(x) -> Optional[TensorMetadata]: else: continue if 'val' in n.meta and test_infer: - (n_meta_val, fake_val) = ((n.meta['val'],),(fake_value,)) if not isinstance(n.meta['val'],(Tuple,List)) else (n.meta['val'], fake_value) + (n_meta_val, fake_val) = ((n.meta['val'],),(fake_value,)) if not isinstance(n.meta['val'],(Tuple,List)) else (n.meta['val'], fake_value) for i,(meta_i,fv_i) in enumerate(zip(n_meta_val, fake_val)): + if not isinstance(fv_i, FakeTensor): + continue assert meta_i.size() == fv_i.size(), "check infer size failed" assert meta_i.dtype == fv_i.dtype, "check infer dtype failed" assert meta_i.stride() == fv_i.stride(), "check infer stride failed" diff --git a/dicp/test/op/test_scalar_tensor.py b/dicp/test/op/test_scalar_tensor.py index 4b555e1f1..8ebdf0259 100644 --- a/dicp/test/op/test_scalar_tensor.py +++ b/dicp/test/op/test_scalar_tensor.py @@ -29,7 +29,7 @@ class TestScalarTensor(): @pytest.mark.parametrize("inputs", [1.0, 3.0, 0.0]) @pytest.mark.parametrize("compiled_model", compiled_model) def test_torch_scalar_tensor(self, inputs, dtype, compiled_model): - redundant_input = torch.randn(1, dtype=dtype) + redundant_input = torch.ones(1, dtype=dtype) _, output = model(inputs, redundant_input, dtype) dynamo.reset() update_dynamo_config(compiled_model.dynamic) From 013ff0c9005dc7242290f606d64d6a254b01b1b4 Mon Sep 17 00:00:00 2001 From: tangzhiyi Date: Fri, 26 Jan 2024 02:14:24 +0000 Subject: [PATCH 3/7] update graph.py --- dicp/dicp/dynamo_bridge/graph.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/dicp/dicp/dynamo_bridge/graph.py b/dicp/dicp/dynamo_bridge/graph.py index fd338c7c3..8b94d1443 100644 --- a/dicp/dicp/dynamo_bridge/graph.py +++ b/dicp/dicp/dynamo_bridge/graph.py @@ -65,10 +65,11 @@ def make_tensor_meta(x) -> Optional[TensorMetadata]: for i,(meta_i,fv_i) in enumerate(zip(n_meta_val, fake_val)): if not isinstance(fv_i, FakeTensor): continue - assert meta_i.size() == fv_i.size(), "check infer size failed" - assert meta_i.dtype == fv_i.dtype, "check infer dtype failed" - assert meta_i.stride() == fv_i.stride(), "check infer stride failed" - assert meta_i.storage_offset() == fv_i.storage_offset(), "check infer storage offset failed" + log_info = f"target: {n.target}, meta_i: {meta_i}, fv_i: {fv_i}" + assert meta_i.size() == fv_i.size(), f"check infer size failed, {log_info}" + assert meta_i.dtype == fv_i.dtype, f"check infer dtype failed, {log_info}" + assert meta_i.stride() == fv_i.stride(), f"check infer stride failed, {log_info}" + assert meta_i.storage_offset() == fv_i.storage_offset(), f"check infer storage offset failed, {log_info}" if 'val' not in n.meta: n.meta['val'] = fake_value n.meta["tensor_meta"] = make_tensor_meta(n.meta['val']) From 9a4665e559f66f37454c8b117b85e65544e28be8 Mon Sep 17 00:00:00 2001 From: tangzhiyi Date: Fri, 26 Jan 2024 02:30:23 +0000 Subject: [PATCH 4/7] update test_scalar_tensor --- dicp/dicp/vendor/AscendGraph/opset_convert.py | 1 + dicp/test/op/test_scalar_tensor.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/dicp/dicp/vendor/AscendGraph/opset_convert.py b/dicp/dicp/vendor/AscendGraph/opset_convert.py index 39c2ef172..2720d99e6 100644 --- a/dicp/dicp/vendor/AscendGraph/opset_convert.py +++ b/dicp/dicp/vendor/AscendGraph/opset_convert.py @@ -82,6 +82,7 @@ def ascendgraph_opset_convert( if is_torch_210: gm = BackendPatternMatcherTransformer( ascend_pattern_matcher, aten_patterns_cls_list).transform(gm) + gm.print_readable() gm = AtenToAscendTransformer(gm).transform() # For bug in pytorch diff --git a/dicp/test/op/test_scalar_tensor.py b/dicp/test/op/test_scalar_tensor.py index 8ebdf0259..5e39043d8 100644 --- a/dicp/test/op/test_scalar_tensor.py +++ b/dicp/test/op/test_scalar_tensor.py @@ -16,7 +16,8 @@ def forward(self, a, redundant_input, dtype): # To solve this problem, an additional redundant input is added, # and the result of an addition operator is returned. res_default = torch.ops.aten.scalar_tensor.default(a, dtype=dtype) - return res_default + redundant_input, res_default + redundant_output = torch.ops.aten.add.Tensor(redundant_input, res_default) + return redundant_output, res_default model = OpModule() From 5a53da26bedd642fcce0423db93c367f77e27dfd Mon Sep 17 00:00:00 2001 From: tangzhiyi Date: Fri, 26 Jan 2024 03:06:10 +0000 Subject: [PATCH 5/7] fix full op --- dicp/dicp/vendor/AscendGraph/conversion.py | 5 +++-- dicp/dicp/vendor/AscendGraph/opset_convert.py | 1 - 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dicp/dicp/vendor/AscendGraph/conversion.py b/dicp/dicp/vendor/AscendGraph/conversion.py index 76d9aa3b4..ca981d9ab 100644 --- a/dicp/dicp/vendor/AscendGraph/conversion.py +++ b/dicp/dicp/vendor/AscendGraph/conversion.py @@ -600,9 +600,10 @@ def view_as_complex(self, x): @register_conversion(torch.ops.aten.full.default) def full(self, dims, value, dtype=torch.float32, layout=torch.strided, device='cpu', pin_memory=False, memory_format=torch.preserve_format): - if len(dims) == 0: - dims = [1] torch_dtype = dtype if dtype else torch.get_default_dtype() + if len(dims) == 0: + return self.get_const_proxy(value, torch_dtype) + dims = [dim.node.meta['val'] if isinstance(dim, torch.fx.proxy.Proxy) and hasattr( dim.node, 'meta') else dim for dim in dims] if isinstance(value, torch.fx.proxy.Proxy) and hasattr(value.node, 'meta'): diff --git a/dicp/dicp/vendor/AscendGraph/opset_convert.py b/dicp/dicp/vendor/AscendGraph/opset_convert.py index 2720d99e6..39c2ef172 100644 --- a/dicp/dicp/vendor/AscendGraph/opset_convert.py +++ b/dicp/dicp/vendor/AscendGraph/opset_convert.py @@ -82,7 +82,6 @@ def ascendgraph_opset_convert( if is_torch_210: gm = BackendPatternMatcherTransformer( ascend_pattern_matcher, aten_patterns_cls_list).transform(gm) - gm.print_readable() gm = AtenToAscendTransformer(gm).transform() # For bug in pytorch From 9a6e0b10abd171feda6d933f124cebe536acb418 Mon Sep 17 00:00:00 2001 From: tangzhiyi Date: Mon, 29 Jan 2024 02:17:58 +0000 Subject: [PATCH 6/7] update --- dicp/test/op/test_scalar_tensor.py | 31 +++++++++++++++++++++++++----- 1 file changed, 26 insertions(+), 5 deletions(-) diff --git a/dicp/test/op/test_scalar_tensor.py b/dicp/test/op/test_scalar_tensor.py index 5e39043d8..05894c24c 100644 --- a/dicp/test/op/test_scalar_tensor.py +++ b/dicp/test/op/test_scalar_tensor.py @@ -9,31 +9,52 @@ class OpModule(torch.nn.Module): + def forward(self, a, dtype): + res_default = torch.ops.aten.scalar_tensor.default(a, dtype=dtype) + return res_default + + +class AscendOpModule(torch.nn.Module): def forward(self, a, redundant_input, dtype): # If there is only one operator called scalar_tensor, # ascend graph compiler will give an error: # GE.. [Check][Param] SetInputs failed: input operator size can not be 0. # To solve this problem, an additional redundant input is added, # and the result of an addition operator is returned. - res_default = torch.ops.aten.scalar_tensor.default(a, dtype=dtype) - redundant_output = torch.ops.aten.add.Tensor(redundant_input, res_default) - return redundant_output, res_default + scalar_tensor = torch.ops.aten.scalar_tensor.default(a, dtype=dtype) + res_default = torch.ops.aten.add.Tensor(redundant_input, scalar_tensor) + return scalar_tensor, res_default model = OpModule() +ascend_model = AscendOpModule() args = parse_args() compiled_model = compile_model(model, args.backend, args.dynamic) +ascend_compiled_model = compile_model(ascend_model, args.backend, args.dynamic) class TestScalarTensor(): + @pytest.mark.skipif(args.backend == 'ascendgraph', reason="skip ascendgraph") @pytest.mark.parametrize("dtype", [torch.float32, torch.int64, torch.float16]) @pytest.mark.parametrize("inputs", [1.0, 3.0, 0.0]) @pytest.mark.parametrize("compiled_model", compiled_model) def test_torch_scalar_tensor(self, inputs, dtype, compiled_model): + output = model(inputs, dtype) + dynamo.reset() + update_dynamo_config(compiled_model.dynamic) + dicp_output = compiled_model.model(inputs, dtype) + + assert torch.allclose(output, dicp_output.cpu(), equal_nan=True) + + @pytest.mark.skipif(args.backend == 'topsgraph', reason="skip topsgraph") + @pytest.mark.parametrize("dtype", [torch.float32, torch.int64, torch.float16]) + @pytest.mark.parametrize("inputs", [1.0, 2.0, 3.0]) + @pytest.mark.parametrize("compiled_model", ascend_compiled_model) + def test_torch_ascend_scalar_tensor(self, inputs, dtype, compiled_model): redundant_input = torch.ones(1, dtype=dtype) - _, output = model(inputs, redundant_input, dtype) + output, _ = ascend_model(inputs, redundant_input, dtype) dynamo.reset() update_dynamo_config(compiled_model.dynamic) - _, dicp_output = compiled_model.model(inputs, redundant_input, dtype) + dicp_output, _ = compiled_model.model(inputs, redundant_input, dtype) assert torch.allclose(output, dicp_output.cpu(), equal_nan=True) From 2b01160d207f6a25b861d0260e71bfffb83d0573 Mon Sep 17 00:00:00 2001 From: tangzhiyi Date: Mon, 29 Jan 2024 03:24:11 +0000 Subject: [PATCH 7/7] update --- dicp/dicp/vendor/AscendGraph/conversion.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/dicp/dicp/vendor/AscendGraph/conversion.py b/dicp/dicp/vendor/AscendGraph/conversion.py index ca981d9ab..c662c33ea 100644 --- a/dicp/dicp/vendor/AscendGraph/conversion.py +++ b/dicp/dicp/vendor/AscendGraph/conversion.py @@ -601,6 +601,8 @@ def view_as_complex(self, x): def full(self, dims, value, dtype=torch.float32, layout=torch.strided, device='cpu', pin_memory=False, memory_format=torch.preserve_format): torch_dtype = dtype if dtype else torch.get_default_dtype() + # If len(dims) == 0, it means this is a scalar tensor with a dimension of 0, + # and it can directly return a const node to construct a scalar tensor. if len(dims) == 0: return self.get_const_proxy(value, torch_dtype)