594 lines
19 KiB
Python
594 lines
19 KiB
Python
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import unittest
|
|
|
|
from test_case_base import (
|
|
TestCaseBase,
|
|
test_instruction_translator_cache_context,
|
|
)
|
|
|
|
import paddle
|
|
from paddle.jit.sot.psdb import check_no_breakgraph
|
|
from paddle.jit.sot.utils import (
|
|
ConditionalFallbackError,
|
|
allow_dynamic_shape_guard,
|
|
enable_0_size_fallback_guard,
|
|
specialized_dim_numbers_guard,
|
|
)
|
|
|
|
|
|
def dynamic_shape_input_func1(x):
|
|
s = x.shape[0]
|
|
return x + s
|
|
|
|
|
|
def dynamic_int_input_func1(x, n):
|
|
x = paddle.reshape(x, [n, -1])
|
|
return (x + n) * 2 - 1, (-n + 1) * 2 - 1, type(n) is int
|
|
|
|
|
|
def dynamic_shape_with_constraints(x, n):
|
|
return (x + n) * 2
|
|
|
|
|
|
def dynamic_int_input_func2(x, n):
|
|
return x + n[1]
|
|
|
|
|
|
def dynamic_int_input_func3(x, n):
|
|
if n < 4:
|
|
return 1
|
|
x = paddle.reshape(x, [n, -1])
|
|
return (x + n) * 2 - 1, (-n + 1) * 2 - 1
|
|
|
|
|
|
def dynamic_shape_access_inner_var_shape(x):
|
|
y = x + 1
|
|
return y.shape[0]
|
|
|
|
|
|
def dynamic_shape_in_list(x, shape):
|
|
return x.reshape(shape)
|
|
|
|
|
|
def dynamic_shape_int_mul_float(x):
|
|
y = x * 0.5
|
|
z = math.sin(y) # Trigger get_py_value
|
|
return z
|
|
|
|
|
|
def dynamic_shape_constraint(x):
|
|
s0, s1, *_ = x.shape
|
|
if s0 < 5:
|
|
return s0 + x
|
|
elif s0 < s1:
|
|
return s0 + x + 1
|
|
elif 2 * (s0 + s1 - 2) <= 30:
|
|
return s0 + x + 2
|
|
else:
|
|
return s0 + x + 3
|
|
|
|
|
|
class CustomConv(paddle.nn.Conv2D):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
@paddle.jit.to_static(full_graph=False)
|
|
def forward(self, x):
|
|
return paddle.nn.functional.conv2d(
|
|
x,
|
|
self.weight,
|
|
self.bias,
|
|
[self._stride[0] + 1, self._stride[1]],
|
|
self._padding,
|
|
self._dilation,
|
|
self._groups,
|
|
self._data_format,
|
|
)
|
|
|
|
|
|
def pool2d_fallback(x, kernel_size):
|
|
return paddle.nn.functional.max_pool2d(x, kernel_size=kernel_size)
|
|
|
|
|
|
class TestOpcodeExecutorDynamicShapeCache(TestCaseBase):
|
|
def test_dynamic_int_input_cache_hit_case1(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
self.assert_results(
|
|
dynamic_int_input_func1, paddle.randn([4, 5, 6]), 2
|
|
)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
for i in range(3, 7):
|
|
self.assert_results(
|
|
dynamic_int_input_func1, paddle.randn([4, 5, 6]), i
|
|
)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
|
|
def test_dynamic_int_input_cache_hit_case2(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
self.assert_results(
|
|
dynamic_int_input_func2, paddle.randn([4, 5, 6]), {1: 2}
|
|
)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
for i in range(3, 7):
|
|
self.assert_results(
|
|
dynamic_int_input_func2, paddle.randn([4, 5, 6]), {1: i}
|
|
)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
|
|
def test_dynamic_int_input_cache_hit_case3(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
translate_count_map = {
|
|
0: 1,
|
|
1: 2, # 1 is dynamic dim
|
|
2: 2, # 2 hit cache, no recompile
|
|
3: 2, # 3 hit cache, no recompile
|
|
4: 3, # 4 is dynamic dim, but it not hit cache
|
|
5: 3, # 5 hit cache, no recompile
|
|
}
|
|
for i in range(0, 6):
|
|
self.assert_results(
|
|
dynamic_int_input_func3, paddle.randn([4, 5, 6]), i
|
|
)
|
|
self.assertEqual(ctx.translate_count, translate_count_map[i])
|
|
|
|
def test_dynamic_shape_input_cache_hit_case1(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
self.assert_results(
|
|
dynamic_shape_input_func1, paddle.randn([2, 4, 5])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
for i in range(3, 7):
|
|
self.assert_results(
|
|
dynamic_shape_input_func1, paddle.randn([i, 4, 5])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
|
|
def test_dynamic_shape_input_cache_hit_case2(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
self.assert_results(
|
|
dynamic_shape_access_inner_var_shape, paddle.randn([2, 4, 5])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
for i in range(3, 7):
|
|
self.assert_results(
|
|
dynamic_shape_access_inner_var_shape,
|
|
paddle.randn([i, 4, 5]),
|
|
)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
|
|
def test_dynamic_shape_cast(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
func1 = check_no_breakgraph(lambda n: bool(n))
|
|
# TODO(SigureMo): Open these cases
|
|
# func2 = check_no_breakgraph(lambda n: int(n))
|
|
# func3 = check_no_breakgraph(lambda n: float(n))
|
|
for func in [func1]:
|
|
self.assert_results(func, 1)
|
|
self.assert_results(func, 2)
|
|
|
|
def test_dynamic_shape_in_list(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
self.assert_results(
|
|
dynamic_shape_in_list,
|
|
paddle.randn([2, 2, 5]),
|
|
[4, 5],
|
|
)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
for i in range(3, 7):
|
|
self.assert_results(
|
|
dynamic_shape_in_list,
|
|
paddle.randn([i, 2, 5]),
|
|
[i * 2, 5],
|
|
)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
|
|
def test_conv_dynamic_shape_stride_fallback(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
for i in range(1, 5):
|
|
conv = CustomConv(3, 3, 3, stride=i)
|
|
conv(paddle.randn([1, 3, 224, 224]))
|
|
self.assertEqual(ctx.translate_count, i)
|
|
|
|
def test_conv_dynamic_shape_kernel_size_fallback(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
for i in range(1, 5):
|
|
x = paddle.randn([1, 3, 224, 224])
|
|
self.assert_results(pool2d_fallback, x, i)
|
|
self.assertEqual(ctx.translate_count, i)
|
|
|
|
def test_pad_dynamic_shape_fallback(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
pad_func = check_no_breakgraph(
|
|
lambda x, n: paddle.nn.functional.pad(x, [0, n, 0, 0])
|
|
)
|
|
for i in range(1, 5):
|
|
self.assert_results(pad_func, paddle.randn([1, 3, 224, 224]), i)
|
|
self.assertEqual(ctx.translate_count, 1 if i == 1 else 2)
|
|
|
|
def test_dynamic_shape_int_mul_float(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
for i in range(1, 6):
|
|
self.assert_results(dynamic_shape_int_mul_float, i)
|
|
|
|
def test_dynamic_shape_constraint(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
const_dim = 6
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([1, 1, const_dim])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([2, 2, const_dim])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 2) # add constraint s0 < 5
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([3, 3, const_dim])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 2) # hit constraint s0 < 5
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([4, 4, const_dim])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 2) # hit constraint s0 < 5
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([5, 6, const_dim])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 3) # add constraint s0 < s1
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([6, 7, const_dim])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 3) # hit constraint s0 < s1
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([7, 8, const_dim])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 3) # hit constraint s0 < s1
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([8, 7, const_dim])
|
|
)
|
|
self.assertEqual(
|
|
ctx.translate_count,
|
|
4, # add constraint 2 * (s0 + s1 - 2) <= 30
|
|
)
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([9, 8, const_dim])
|
|
)
|
|
self.assertEqual(
|
|
ctx.translate_count,
|
|
4, # hit constraint 2 * (s0 + s1 - 2) <= 30
|
|
)
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([10, 9, const_dim])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 5) # add constraint else
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([11, 10, const_dim])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 5) # hit constraint else
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([4, 3, const_dim])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 5) # hit constraint s0 < 5
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([5, 8, const_dim])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 5) # hit constraint s0 < s1
|
|
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([8, 8, const_dim])
|
|
)
|
|
self.assertEqual(
|
|
ctx.translate_count,
|
|
5, # hit 2 * (s0 + s1 - 2) <= 30
|
|
)
|
|
|
|
with self.assertRaises(ConditionalFallbackError):
|
|
self.assert_results(
|
|
dynamic_shape_constraint, paddle.randn([0, 1, const_dim])
|
|
)
|
|
|
|
def test_mixed_dynamic_and_static(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
a = paddle.randn([4, 5, 6])
|
|
self.assert_results(dynamic_int_input_func1, a, 1)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
self.assert_results(dynamic_int_input_func1, a, 0)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
for i in range(2, 6):
|
|
self.assert_results(dynamic_int_input_func1, a, i)
|
|
self.assertEqual(ctx.translate_count, 3)
|
|
|
|
def test_mixed_static_after_dynamic(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
a = paddle.randn([4, 5, 6])
|
|
self.assert_results(dynamic_int_input_func1, a, 2)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
for i in range(3, 6):
|
|
self.assert_results(dynamic_int_input_func1, a, i)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
self.assert_results(dynamic_int_input_func1, a, 0)
|
|
self.assertEqual(ctx.translate_count, 3)
|
|
self.assert_results(dynamic_int_input_func1, a, 1)
|
|
self.assertEqual(ctx.translate_count, 3)
|
|
|
|
def test_dynamic_shape_with_constraints(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
self.assert_results(
|
|
dynamic_shape_with_constraints, paddle.randn([4, 5, 6]), 2
|
|
)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
for i in range(3, 7):
|
|
self.assert_results(
|
|
dynamic_shape_with_constraints,
|
|
paddle.randn([4 + i, 5, 6]),
|
|
i,
|
|
)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
|
|
|
|
@check_no_breakgraph
|
|
def dynamic_shape_non_break_non_inplace_ops(x):
|
|
s0 = x.shape[0]
|
|
s1 = s0 + 1
|
|
s2 = 1 + s0
|
|
s3 = s1 + s2
|
|
s4 = s1 * s2
|
|
s5 = s1 - s2
|
|
s6 = s1 / s2
|
|
s7 = s1 // s2
|
|
s8 = s1 % s2
|
|
s9 = s1**s2
|
|
s10 = s1 & s2
|
|
s11 = s1 | s2
|
|
s12 = s1 ^ s2
|
|
s13 = s1 << s2
|
|
s14 = s1 >> s2
|
|
s15 = s1 == s2
|
|
s16 = s1 != s2
|
|
s17 = s1 < s2
|
|
s18 = s1 <= s2
|
|
s19 = s1 > s2
|
|
s20 = s1 >= s2
|
|
s21 = bool(s1)
|
|
s22 = not s2
|
|
return (
|
|
s0,
|
|
s1,
|
|
s2,
|
|
s3,
|
|
s4,
|
|
s5,
|
|
s6,
|
|
s7,
|
|
s8,
|
|
s9,
|
|
s10,
|
|
s11,
|
|
s12,
|
|
s13,
|
|
s14,
|
|
s15,
|
|
s16,
|
|
s17,
|
|
s18,
|
|
s19,
|
|
s20,
|
|
s21,
|
|
s22,
|
|
)
|
|
|
|
|
|
@check_no_breakgraph
|
|
def dynamic_shape_non_break_inplace_ops(x):
|
|
s0 = x.shape[0]
|
|
s1 = s0 + 1
|
|
s2 = s3 = s4 = s5 = s6 = s7 = s8 = s9 = s10 = s11 = s12 = s13 = s0
|
|
s2 += s1
|
|
s3 *= s1
|
|
s4 -= s1
|
|
# TODO(SigureMo): Open this case, currently the compute result between Python and C++ (Paddle Kernel)
|
|
# has a small difference (0.8333333134651184 and 0.8333333333333334)
|
|
# s5 /= s1
|
|
s6 //= s1
|
|
s7 %= s1
|
|
s8 **= s1
|
|
s9 &= s1
|
|
s10 |= s1
|
|
s11 ^= s1
|
|
s12 <<= s1
|
|
s13 >>= s1
|
|
return (
|
|
s0,
|
|
s1,
|
|
s2,
|
|
s3,
|
|
s4,
|
|
# s5,
|
|
s6,
|
|
s7,
|
|
s8,
|
|
s9,
|
|
s10,
|
|
s11,
|
|
s12,
|
|
s13,
|
|
)
|
|
|
|
|
|
class TestDynamicShapeNonBreakOps(TestCaseBase):
|
|
def test_dynamic_shape_non_break_non_inplace_ops(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
self.assert_results(
|
|
dynamic_shape_non_break_non_inplace_ops, paddle.randn([4, 5, 6])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
for i in range(5, 9):
|
|
self.assert_results(
|
|
dynamic_shape_non_break_non_inplace_ops,
|
|
paddle.randn([i, 5, 6]),
|
|
)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
|
|
def test_dynamic_shape_non_break_inplace_ops(self):
|
|
with (
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
):
|
|
self.assert_results(
|
|
dynamic_shape_non_break_inplace_ops, paddle.randn([4, 5, 6])
|
|
)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
for i in range(5, 9):
|
|
self.assert_results(
|
|
dynamic_shape_non_break_inplace_ops,
|
|
paddle.randn([i, 5, 6]),
|
|
)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
|
|
|
|
def dynamic_shape_for_specialized_dim_numbers(x):
|
|
return x + 1
|
|
|
|
|
|
class TestSpecializedDimNumbers(TestCaseBase):
|
|
def test_specialized_dim_numbers_01(self):
|
|
with (
|
|
specialized_dim_numbers_guard("01"),
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
enable_0_size_fallback_guard(False),
|
|
):
|
|
x = paddle.randn([0, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
x = paddle.randn([1, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
x = paddle.randn([2, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 3)
|
|
x = paddle.randn([3, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 3)
|
|
|
|
def test_specialized_dim_numbers_0(self):
|
|
with (
|
|
specialized_dim_numbers_guard("0"),
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
enable_0_size_fallback_guard(False),
|
|
):
|
|
x = paddle.randn([0, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
x = paddle.randn([1, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
x = paddle.randn([2, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
x = paddle.randn([3, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
|
|
def test_specialized_dim_numbers_no(self):
|
|
with (
|
|
specialized_dim_numbers_guard("no"),
|
|
allow_dynamic_shape_guard(True),
|
|
test_instruction_translator_cache_context() as ctx,
|
|
enable_0_size_fallback_guard(False),
|
|
):
|
|
x = paddle.randn([10, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 1)
|
|
x = paddle.randn([0, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
x = paddle.randn([1, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
x = paddle.randn([2, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
x = paddle.randn([3, 5, 6])
|
|
self.assert_results(dynamic_shape_for_specialized_dim_numbers, x)
|
|
self.assertEqual(ctx.translate_count, 2)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|