وضعیت فعلی پویایی به طور رسمی‌تر در RFC مربوط به پویایی شرح داده شده است، این صفحه یک مرور کلی سطح بالا از RFC ارائه می‌دهد و APIها و ابزارهای مهم برای تعامل با برنامه‌های پویا را مورد بحث قرار می‌دهد.
اصطلاحات و پشتیبانی دینامیک (Dynamism) به صورت اجمالی
ابتدا، برای پوشش چند اصطلاح که در این سند ظاهر می‌شوند، و همچنین مقدمه‌ای کوتاه در مورد پشتیبانی آنها در StableHLO:
ابعاد پویا
ابعاد پویا به هر بعدی اشاره دارد که اندازه بُعد آن ناشناخته است. در StableHLO ما ابعاد پویا را با استفاده از ? نشان می‌دهیم، یعنی tensor<16x?xf32> .
پویایی محدود
پویایی کراندار به یک بُعد پویا اشاره دارد که مقدار آن دارای یک حد بالای شناخته شده است. عموماً این برای پر کردن تانسور در حین اجرا مفید است. در StableHLO، ما پویایی کراندار را با استفاده از #stablehlo.bounds به عنوان یک کدگذاری تانسور نشان می‌دهیم، یعنی یک تانسور رتبه ۲ با یک بُعد پویا که در ۱۶ کراندار شده و بُعد دیگر بدون کران می‌تواند به صورت tensor<?x?xf32, #stablehlo.bounds<16, ?>> نمایش داده شود.
StableHLO قادر به نمایش پویایی محدود است، اما پشتیبانی چارچوب محدودی دارد که منشأ آن TensorFlow است و با مقداری پشتیبانی در PyTorch/XLA ارائه می‌شود.
پویایی بی‌حد و مرز
پویایی نامحدود، همانطور که از نامش پیداست، به یک بُعد پویا بدون هیچ حد مشخصی برای اندازه اشاره دارد. این نوع پویایی در StableHLO بسیار رایج است و از JAX، PyTorch/XLA و TF پشتیبانی می‌کند و اغلب برای خروجی گرفتن از مدل‌هایی با اندازه دسته‌ای پویا یا طول توالی استفاده می‌شود.
در StableHLO ما به سادگی کدگذاری کرانه‌ها را برای این نوع پویایی، یعنی tensor<?x?xf32> حذف می‌کنیم.
چندریختی شکل
چندریختی شکل اصطلاحی است که از JAX به ارث برده‌ایم .
دو پیامد کلیدی برای شکل‌دهی چندریختی وجود دارد:
- تمام پویایی برنامه به آرگومان‌های ورودی آن برمی‌گردد.
- تمام پویایی فقط مربوط به شکل‌های تانسور است، یعنی وابسته به داده‌ها نیست.
با این دو قانون، وقتی شکل‌های استاتیک یک برنامه مشخص شد، می‌توانیم یک برنامه پویا را گرفته و آن را به طور کامل به یک برنامه استاتیک برای کامپایل تبدیل کنیم (به «گذرگاه‌های کامپایلر برای پالایش برنامه‌های پویا» مراجعه کنید).
عموماً چندریختی شکل از پویایی نامحدود استفاده می‌کند، اگر شکل‌های آرگومان شناخته شده بتوانند به یک برنامه کاملاً ایستا منجر شوند، نیازی به حدس زدن در مورد نحوه محدود کردن مقادیر نیست.
پویایی وابسته به داده
پویایی وابسته به داده به ابعاد پویا و اندازه‌هایی اشاره دارد که مربوط به داده‌های درون یک تانسور هستند. مثال متعارف آن یک تابع nonzeros است که اندیس تمام عناصری را که در یک مقدار تانسور برابر 0 هستند، برمی‌گرداند. شکل را نمی‌توان بدون ارزیابی داده‌ها شناخت، اما اغلب می‌توان آن را با استفاده از پویایی محدود کامپایل کرد و حافظه اضافی را صرف اندازه تانسور خروجی بالقوه کرد.
بسیاری از عملیات پویای وابسته به داده را می‌توان با استفاده از پویایی محدود مدل‌سازی کرد، که در آن یک حد بالا برای اندازه تانسور مشخص شده است و سخت‌افزار عموماً این را از طریق پدینگ تانسور پیاده‌سازی می‌کند. امروزه در PyTorch/XLA و TensorFlow پشتیبانی از پویایی وابسته به داده وجود دارد، اما JAX در حال حاضر عملیاتی را که منجر به پویایی وابسته به داده می‌شوند، ردیابی نمی‌کند.
خروجی گرفتن از برنامه‌ها با ابعاد پویا
برای کسب اطلاعات در مورد نحوه خروجی گرفتن از برنامه‌ها با اندازه‌های دسته‌ای پویا یا طول توالی، به آموزش‌های StableHLO ما مراجعه کنید:
- آموزش JAX > خروجی گرفتن با اندازه دسته‌ای پویا
- آموزش PyTorch/XLA > خروجی گرفتن با اندازه دسته‌ای پویا
کامپایلر برای اصلاح برنامه‌های پویا، مجوزهایی صادر می‌کند.
خط لوله عبور دینامیک را حذف کنید
چند مسیر مفید برای اصلاح شکل‌ها وجود دارد که به راحتی همه آنها در یک مسیر به نام createStablehloRemoveDynamismPipeline قرار گرفته‌اند:
void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
TypeRange refinedTypes);
پاس‌های انفرادی برای بهبود پویایی
به طور جداگانه، مسیرهایی که برای اصلاح شکل مفید هستند عبارتند از:
-
stablehlo-refine-argumentsبرای جایگزینی آرگومان‌های ورودی با انواع تانسورهای مشخص. -
stablehlo-refine-shapesبرای انتشار اطلاعات شکل آرگومان ورودی جدید در کل برنامه. -
stablehlo-canonicalize-dynamismبرای جایگزینی عملیات پویا با انواع استاتیک آنها. -
stablehlo-check-shape-assertionsبرای بررسی و حذف فراخوانی‌های سفارشیِ shape assertions.
برای اطلاعات به‌روز و مثال‌ها، به مستندات مرتبط مراجعه کنید.
مثال: پویایی چه فایده‌ای دارد و چگونه می‌توانم از آن استفاده کنم؟
پویایی کاربردهای زیادی دارد، در اینجا ما عمدتاً بر روی مورد استفاده رایج برای چندریختی شکل تمرکز خواهیم کرد - ایجاد یک نمایش مدل خروجی انعطاف‌پذیر، که عموماً برای نمایش اندازه دسته پویا یا طول توالی استفاده می‌شود.
مدل استاتیک add_one
ما از مدل ساده‌ی add_one زیر برای نمایش این موضوع استفاده خواهیم کرد:
def add_one(x):
return x + 1
وقتی با استفاده از tensor<4xf32> ردیابی انجام شود، برنامه StableHLO زیر را خواهیم داشت:
// File: add_one.mlir
func.func @add_one(%arg0: tensor<4xf32>) -> tensor<4xf32> {
%cst = stablehlo.constant dense<1.000000e+00> : tensor<4xf32>
%0 = stablehlo.add %arg0, %cst : tensor<4xf32>
return %0 : tensor<4xf32>
}
این مدل فقط برای آرگومان‌های ورودی که شکل tensor<4xf32> دارند، کار می‌کند. اگر اندازه دسته یا طول توالی خود را تغییر دهیم، باید کد منبع را دوباره ردیابی کنیم و دوباره به StableHLO کاهش دهیم، و هیچ تضمینی وجود ندارد که حتی هنوز به کد منبع دسترسی داشته باشیم!
مدل افزونه پویا
اینجاست که پویایی چندریختی شکل وارد عمل می‌شود. در عوض، JAX و PyTorch/XLA می‌توانند مدل add_one را با IR معتبر پویا منتشر کنند که ثابت را برای مطابقت با شکل ورودی پویا به صورت زیر پخش می‌کند:
// File: add_one_dynamic.mlir
func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
%cst = stablehlo.constant dense<1.0> : tensor<f32>
%0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor<?xf32>) -> tensor<i32>
%1 = stablehlo.reshape %0 : (tensor<i32>) -> tensor<1xi32>
%2 = stablehlo.dynamic_broadcast_in_dim %cst, %1, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
%3 = stablehlo.add %arg0, %2 : tensor<?xf32>
return %3 : tensor<?xf32>
}
این نمایش مدل بسیار انعطاف‌پذیرتر است و امکان تعیین مقادیری مانند اندازه دسته یا طول توالی را با تأخیر فراهم می‌کند. این مدل را می‌توان روی پلتفرم‌هایی با پشتیبانی از شکل پویا (مانند AI Edge ) مستقر کرد، یا می‌توان آن را با استفاده از مسیرهای پویایی که در این مستندات ذکر شده است، اصلاح کرد.
اصلاح مدل پویا
برای مثال، ترتیب عبور زیر می‌تواند این برنامه را به طور کامل اصلاح کند:
stablehlo-opt add_one_dynamic.mlir \
--stablehlo-refine-arguments='types=tensor<16xf32>' \
--stablehlo-refine-shapes \
--stablehlo-canonicalize-dynamism
به تدریج، برنامه به این صورت تغییر می‌کند:
// After stablehlo-refine-arguments: Inputs updated, shapes not propagated
func.func public @main(%arg0: tensor<16xf32>) -> tensor<?xf32> {
%c = stablehlo.constant dense<16> : tensor<1xi64>
%0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {indices_of_shape_operands = dense<1> : tensor<1xi64>} : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
...
%3 = stablehlo.dynamic_broadcast_in_dim %cst, %2, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
%4 = stablehlo.add %0, %3 : tensor<?xf32>
return %4 : tensor<?xf32>
}
// After stablehlo-refine-shapes: Shapes propagated, dynamic ops still exist
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
%cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
%c = stablehlo.constant dense<16> : tensor<1xi32>
%0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>
%1 = stablehlo.add %arg0, %0 : tensor<16xf32>
return %1 : tensor<16xf32>
}
// After stablehlo-canonicalize-dynamism: Dynamic ops replaced with static ops
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
%cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
%0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>
%1 = stablehlo.add %arg0, %0 : tensor<16xf32>
return %1 : tensor<16xf32>
}
// (Bonus) Use ` --stablehlo-aggressive-simplification` pass to canonicalize the
// constant broadcast, leaving us with the original static program in this case.
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
%cst = stablehlo.constant dense<1.000000e+00> : tensor<16xf32>
%0 = stablehlo.add %arg0, %cst : tensor<16xf32>
return %0 : tensor<16xf32>
}