StableHLO মধ্যে গতিশীলতা

গতিশীলতার বর্তমান অবস্থাটি ডায়নামিজম RFC- তে আরও আনুষ্ঠানিকভাবে বর্ণিত হয়েছে, এই পৃষ্ঠাটি RFC-এর একটি উচ্চ স্তরের ওভারভিউ প্রদান করবে এবং গতিশীল প্রোগ্রামগুলির সাথে ইন্টারঅ্যাক্ট করার জন্য গুরুত্বপূর্ণ API এবং টুলিং নিয়ে আলোচনা করবে।

গতিশীলতার পরিভাষা এবং সহায়তার সংক্ষিপ্তসার

First, to cover a few terms that will appear in this doc, as well as a brief intro to their support in StableHLO:

গতিশীল মাত্রা

Dynamic dimensions refers to any dimension whose dimension size is unknown. In StableHLO we represent dynamic dimensions using ? , ie tensor<16x?xf32> .

সীমাবদ্ধ গতিশীলতা

বাউন্ডেড ডাইনামিজম বলতে এমন একটি ডাইনামিক ডাইমেনশনকে বোঝায় যার মান একটি পরিচিত উপরের সীমানা ধারণ করে। সাধারণত এটি এক্সিকিউশনের সময় টেনসর প্যাড করার জন্য কার্যকর। StableHLO তে আমরা #stablehlo.bounds টেনসর এনকোডিং হিসেবে ব্যবহার করে বাউন্ডেড ডাইনামিজম উপস্থাপন করি, অর্থাৎ একটি র‍্যাঙ্ক-২ টেনসর যার একটি ডাইনামিক ডাইমেনশন 16 এ আবদ্ধ এবং অন্যটি বাউন্ড ছাড়াই tensor<?x?xf32, #stablehlo.bounds<16, ?>> হিসেবে উপস্থাপন করা যেতে পারে।

StableHLO is able to represent bounded dynamism, but there is limited framework support, originating in TensorFlow, and with some support in PyTorch/XLA.

অসীম গতিশীলতা

নাম থেকেই বোঝা যায়, আনবাউন্ডেড ডাইনামিসম বলতে এমন একটি গতিশীল মাত্রা বোঝায় যার আকারের কোনও নির্দিষ্ট সীমা নেই। এই ধরণের ডাইনামিসম StableHLO তে খুবই সাধারণ, যেখানে JAX, PyTorch/XLA এবং TF সাপোর্ট রয়েছে, যা প্রায়শই ডাইনামিক ব্যাচ সাইজ বা সিকোয়েন্স লেন্থ সহ মডেল রপ্তানি করার জন্য ব্যবহৃত হয়।

In StableHLO we simply elide the bounds encoding for this form of dynamism, ie tensor<?x?xf32> .

আকৃতির বহুরূপতা

Shape polymorphism is a term we've inherited from JAX .

There are two key implications to shape polymorphism:

  1. All dynamism in the program traces back to its input arguments.
  2. All dynamism pertains to tensor shapes only, ie not data-dependent.

এই দুটি নিয়মের সাহায্যে, একবার একটি প্রোগ্রামের স্ট্যাটিক আকার জানা গেলে, আমরা একটি ডায়নামিক প্রোগ্রাম নিতে এবং এটিকে সম্পূর্ণরূপে সংকলনের জন্য একটি স্ট্যাটিক প্রোগ্রামে রূপান্তর করতে সক্ষম হই ( "ডাইনামিক প্রোগ্রামগুলি সংশোধন করার জন্য কম্পাইলার পাস" দেখুন)।

Generally shape polymorphism uses unbounded dynamism, if known argument shapes can lead to a fully static program, there isn't a need to guess on how to bound the values.

তথ্য-নির্ভর গতিশীলতা

ডেটা-নির্ভর গতিশীলতা বলতে টেনসরের ভিতরে থাকা ডেটার সাথে সম্পর্কিত গতিশীল মাত্রার আকার বোঝায়। ক্যানোনিকাল উদাহরণ হল একটি nonzeros ফাংশন যা টেনসর মানের 0 থাকা সমস্ত উপাদানের সূচকগুলি ফেরত দেয়। ডেটা মূল্যায়ন না করে আকৃতি জানা যায় না, তবে এটি প্রায়শই বাউন্ডেড গতিশীলতা ব্যবহার করে কম্পাইল করা যেতে পারে, সম্ভাব্য আউটপুট টেনসর আকারের উপর অতিরিক্ত মেমরি ব্যয় করে।

অনেক ডেটা-নির্ভর গতিশীল অপারেশনগুলিকে বাউন্ডেড ডায়নামিজম ব্যবহার করে মডেল করা যেতে পারে, যেখানে টেনসর আকারের উপর একটি উপরের বাউন্ড নির্দিষ্ট করা থাকে এবং হার্ডওয়্যার সাধারণত টেনসর প্যাডিংয়ের মাধ্যমে এটি বাস্তবায়ন করবে। আজ PyTorch/XLA এবং TensorFlow-তে ডেটা-নির্ভর গতিশীলতার জন্য কিছু সমর্থন রয়েছে, কিন্তু JAX বর্তমানে এমন ক্রিয়াকলাপগুলি ট্রেস করে না যা ডেটা-নির্ভর গতিশীলতার দিকে পরিচালিত করে।

গতিশীল মাত্রা সহ প্রোগ্রাম রপ্তানি করা হচ্ছে

See our StableHLO tutorials for information on how to export programs with dynamic batch sizes or sequence lengths:

গতিশীল প্রোগ্রামগুলি পরিমার্জনের জন্য কম্পাইলার পাস

গতিশীলতা পাস পাইপলাইন সরান

There are a few useful passes for refining shapes, conveniently they are all bundled in a pass pipeline createStablehloRemoveDynamismPipeline :

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

গতিশীলতা বৃদ্ধির জন্য ব্যক্তিগত পাস

Individually, the passes that tend to be useful for shape refinement are:

See linked documentation for up-to-date information and examples.

Example: How is dynamism useful, and how can I use it?

ডায়নামিজমের অনেক ব্যবহার রয়েছে, এখানে আমরা মূলত শেপ পলিমরফিজমের সাধারণ ব্যবহারের ক্ষেত্রে ফোকাস করব - একটি নমনীয় রপ্তানি মডেল উপস্থাপনা তৈরি করা, যা সাধারণত ডায়নামিক ব্যাচের আকার বা সিকোয়েন্স দৈর্ঘ্য উপস্থাপন করতে ব্যবহৃত হয়।

স্ট্যাটিক অ্যাড_ওয়ান মডেল

We'll use the following simple add_one model to demonstrate this:

def add_one(x):
  return x + 1

When traced using a tensor<4xf32> we'll get the following StableHLO program:

// 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 গতিশীলভাবে বৈধ IR সহ add_one মডেল নির্গত করতে পারে যা নিম্নরূপ গতিশীল ইনপুট আকৃতির সাথে মেলে ধ্রুবক সম্প্রচার করবে:

// 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 ), অথবা এই ডকুমেন্টেশনে উল্লিখিত গতিশীলতা পাসগুলি ব্যবহার করে এটি পরিমার্জিত করা যেতে পারে।

গতিশীল মডেলকে পরিমার্জন করা

For example the following pass ordering can fully refine this program:

stablehlo-opt add_one_dynamic.mlir \
  --stablehlo-refine-arguments='types=tensor<16xf32>' \
  --stablehlo-refine-shapes \
  --stablehlo-canonicalize-dynamism

Incrementally, this is how the program gets transformed:

// 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>
}