وضعیت فعلی پویایی به طور رسمی در Dynamism RFC توضیح داده شده است، این صفحه نمای کلی سطح بالایی از RFC ارائه می دهد و API های مهم و ابزارهای تعامل با برنامه های پویا را مورد بحث قرار می دهد.
بررسی اجمالی اصطلاحات و پشتیبانی پویایی
ابتدا، برای پوشش چند اصطلاحی که در این سند ظاهر می شوند، و همچنین معرفی مختصری از پشتیبانی آنها در StableHLO:
ابعاد دینامیک
ابعاد پویا به هر بعدی اطلاق می شود که اندازه ابعاد آن ناشناخته است. در StableHLO ابعاد پویا را با استفاده از ? ، یعنی tensor<16x?xf32> .
پویایی محدود
پویایی محدود به بعد پویایی اشاره دارد که مقدار آن دارای کران بالایی شناخته شده است. به طور کلی این برای بالشتک کردن تانسور در حین اجرا مفید است. در StableHLO ما دینامیسم محدود را با استفاده از #stablehlo.bounds به عنوان رمزگذاری تانسور نشان می‌دهیم، یعنی یک تانسور رتبه-2 با یک بعد دینامیکی محدود شده در 16 و دیگری بدون کران را می‌توان به صورت tensor<?x?xf32, #stablehlo.bounds<16, ?>> نشان داد. 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 ما مراجعه کنید:
پاس های کامپایلر برای پالایش برنامه های پویا
خط لوله عبور دینامیسم را حذف کنید
چند پاس مفید برای اصلاح اشکال وجود دارد، به راحتی همه آنها در یک خط لوله گذر createStablehloRemoveDynamismPipeline می‌شوند.
void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
TypeRange refinedTypes);
پاس های انفرادی برای تقویت پویایی
به طور جداگانه، پاس هایی که برای اصلاح شکل مفید هستند عبارتند از:
-
stablehlo-refine-argumentsبرای جایگزینی آرگومان های ورودی با انواع تانسور بتن. -
stablehlo-refine-shapesبرای انتشار اطلاعات شکل آرگومان ورودی جدید در کل برنامه. -
stablehlo-canonicalize-dynamismبرای جایگزینی عملیات پویا با انواع استاتیک آنها.
برای اطلاعات و نمونه‌های به‌روز به اسناد پیوندی مراجعه کنید.
مثال: پویایی چگونه مفید است و چگونه می توانم از آن استفاده کنم؟
Dynamism کاربردهای زیادی دارد، در اینجا ما عمدتاً بر روی مورد استفاده رایج برای Shape Polymorphism تمرکز می کنیم - ایجاد یک نمایش مدل صادراتی انعطاف پذیر، که عموماً برای نشان دادن اندازه دسته پویا یا طول دنباله استفاده می شود.
مدل 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 پایین بیاوریم، و هیچ تضمینی وجود ندارد که حتی به کد منبع همچنان دسترسی داشته باشیم!
مدل add_one پویا
اینجاست که پویایی چندشکلی شکل ظاهر می شود. در عوض 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>
}