1# Copyright 2024 Huawei Technologies Co., Ltd 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================ 15"""Test the feature with one stage""" 16import pytest 17import numpy 18import types 19from mindspore import Tensor, jit 20from mindspore._c_expression import get_code_extra 21 22 23cfg = { 24 "print_after_all": False, 25 "compile_by_trace": True, 26} 27 28 29@pytest.mark.level0 30@pytest.mark.platform_x86_gpu_training 31@pytest.mark.env_onecard 32def test_code_generator_with_complete_graph(): 33 """ 34 Feature: One stage code generate. 35 Description: Test one stage code generate with complete graph. 36 Expectation: No exception. 37 """ 38 39 @jit(mode="PIJit", jit_config={**cfg, "interpret_captured_code": False}) 40 def graph_test(x, y, *args, z = 1, **kw): 41 a = x + y 42 b = y - z 43 c = x * y * z * a * b 44 return c 45 46 @jit(mode="PIJit", jit_config={**cfg, "interpret_captured_code": True}) 47 def code_test(x, y, *args, z = 1, **kw): 48 a = x + y 49 b = y - z 50 c = x * y * z * a * b 51 return c 52 53 x = Tensor(numpy.zeros((4, 4))) 54 y = Tensor(numpy.random.rand(4, 4)) 55 result = graph_test(x, y) 56 excepted = code_test(x, y) 57 58 graph_phase = get_code_extra(graph_test)["code"].get("phase_", None) 59 non_code = get_code_extra(graph_test)["code"].get("compiled_code_", None) 60 non_phase = get_code_extra(code_test)["code"].get("phase_", None) 61 new_code = get_code_extra(code_test)["code"].get("compiled_code_", None) 62 63 assert (result == excepted).asnumpy().all() and not result.asnumpy().all() 64 assert non_code is None and non_phase is None 65 assert isinstance(graph_phase, str) and isinstance(new_code, types.CodeType) 66 67 68@pytest.mark.level0 69@pytest.mark.platform_x86_gpu_training 70@pytest.mark.env_onecard 71def test_code_generator_with_exception(): 72 """ 73 Feature: One stage code generate. 74 Description: Test one stage code generate with exception code. 75 Expectation: Raise exception. 76 """ 77 78 @jit(mode="PIJit", jit_config = {**cfg, "interpret_captured_code":True}) 79 def code_test(x, y, unknown_func, *args, z=1, **kw): 80 a = x + y 81 b = y - z 82 unknown_func(x, (1,2,3,4)) 83 c = x * y * z * a * b 84 return c 85 86 # here not test graph, one stage has bugs 87 88 x = Tensor(numpy.zeros((4, 4))) 89 y = Tensor(numpy.random.rand(4, 4)) 90 unknown_func = Tensor.shape.__set__ # a function with exception 91 92 msg = None 93 try: 94 z = code_test(x, y, unknown_func=unknown_func) 95 print(z) 96 except Exception as e: 97 msg = str(e) 98 99 non_phase = get_code_extra(code_test)["code"].get("phase_", None) 100 new_code = get_code_extra(code_test)["code"].get("compiled_code_", None) 101 102 assert msg 103 assert non_phase is None and isinstance(new_code, types.CodeType) 104