Async HLO নির্দেশাবলী

  1. HLO-তে অ্যাসিঙ্ক অপারেশন যোগ করা কষ্টসাধ্য (যেমন all-reduce-start এবং all-reduce-done )।
  2. কিছু অ্যাসিঙ্ক্রোনাস ব্যবহারের ক্ষেত্রে শুরু এবং শেষের বিভাজন অপর্যাপ্ত হতে পারে।

প্রথম সীমাবদ্ধতাটি মোকাবেলা করার জন্য, আমরা kAsyncStart , kAsyncUpdate , এবং kAsyncDone নামে নতুন এক সেট অ্যাসিঙ্ক্রোনাস অপকোড প্রবর্তন করার প্রস্তাব করছি। এর মূল উদ্দেশ্য হলো একটি জেনেরিক অ্যাসিঙ্ক্রোনাস অপকোড তৈরি করা, যা যেকোনো HLO ইন্সট্রাকশনকে র‍্যাপ করতে পারবে। যে প্রকৃত অপারেশনটি অ্যাসিঙ্ক্রোনাসভাবে সম্পাদিত হবে, তা একটি কলড কম্পিউটেশন ব্যবহার করে এনকোড করা হবে, যার রুট হিসেবে শুধু ইন্সট্রাকশনটি এবং ইনপুট হিসেবে যেকোনো প্যারামিটার থাকবে। এর ফলে, যেকোনো অ্যাসিঙ্ক্রোনাস অপারেশনের জন্য ইন-ফ্লাইট ইনপুট/আউটপুট বাফার হ্যান্ডলিং এবং অ্যালিয়াসিং শেয়ার করা যাবে। এরপর async-start ইন্সট্রাকশনের আউটপুট শেপ হবে ইনপুট অপারেন্ড, আউটপুট ভ্যালু এবং async-update বা async-done ইন্সট্রাকশনের জন্য প্রয়োজনীয় যেকোনো ইন্টারমিডিয়েট স্টেটের একটি টাপল।

%async_op {
  %param0 = f32[64] parameter(0)
  ROOT %op = f32[32] op(f32[64] %param0), op_specific_attr=foo
}

%async-start = ((f32[64]), f32[32], s32[]) async-start(f32[64] %operand),
                                           calls=%async_op
%async-done = f32[32] async-done(((f32[64]), f32[32], s32[]) %async-start)

উপরের উপস্থাপনায়, শুধুমাত্র async-start একটি কল করা কম্পিউটেশন আছে, কারণ async-done কী করে তা এর অপারেন্ড অনুসরণ করে সংশ্লিষ্ট async-start খুঁজে বের করার মাধ্যমে সহজেই জানা যায়, যা থেকে কল করা কম্পিউটেশনটি পাওয়া যায়।

আরও লক্ষ্য করুন যে, async-start এর আউটপুট টাপলের প্রথম উপাদানটি হলো অপারেন্ডগুলো ধারণকারী একটি টাপল। এই অপারেন্ড টাপলের উপাদানগুলো নিজ নিজ অপারেন্ডের সাথে অ্যালিয়াস করে, ফলে তাদের বাফারগুলো অন্তত async-done নির্দেশনা পর্যন্ত সচল থাকে। একইভাবে, দ্বিতীয় উপাদানটি async-done এর আউটপুটের সাথে অ্যালিয়াস করে, এবং তৃতীয় উপাদানটি হলো কনটেক্সট স্টেট যা অ্যাসিঙ্ক্রোনাস অপারেশনটির হিসাব রাখতে ব্যবহৃত হয়। এই উপস্থাপনাটি স্বাভাবিকভাবেই অ্যাসিঙ্ক্রোনাস অপারেশনের ইনপুট এবং/অথবা আউটপুটে একাধিক টেনসর সমর্থন করে:

%async_op {
  %param0 = f32[64] parameter(0)
  %param1 = f32[64] parameter(1)
  ROOT %op = (f32[32], f32[32]) op(f32[64] %param0, f32[64] %param1),
                                op_specific_attr=foo
}

%async-start = ((f32[64], f32[64]), (f32[32], f32[32]), s32[])
               async-start(f32[64] %operand0, f32[64] %operand1),
               calls=%async_op
%async-done = (f32[32], f32[32]) async-done(%async-start)

এছাড়াও, অপারেশনটিকে আরও শূন্য বা তার বেশি async-update ধাপে বিভক্ত করা যেতে পারে যা মধ্যবর্তী গণনা সম্পাদন করে। async-update নির্দেশনার সাথে ইনপুট/আউটপুট অ্যালিয়াসিং একইভাবে কাজ করে এবং প্রতিটি async-startasync-update নির্দেশনার অবশ্যই একজন ব্যবহারকারী থাকতে হবে যিনি হয় অন্য একটি async-update অথবা একটি async-done

%async_op {
  %param0 = f32[64] parameter(0)
  ROOT %op = f32[32] op(f32[64] %param0), op_specific_attr=foo
}

%async-start = ((f32[64]), f32[32], s32[]) async-start(f32[64] %operand),
                                         calls=%async_op
%async-update0 = ((f32[64]), f32[32], s32[]) async-update(
                           ((f32[64]), f32[32], s32[]) %async-start)
%async-update1 = ((f32[64]), f32[32], s32[]) async-update(
                           ((f32[64]), f32[32], s32[]) %async-update0)
%async-done = f32[32] async-done(((f32[64]), f32[32], s32[]) %async-update1)

সিনট্যাক্স সুগার

HLO পার্সারটি অ্যাসিঙ্ক্রোনাস অপারেশনগুলোকে ফার্স্ট-ক্লাস অপকোডের মতো স্বয়ংক্রিয়ভাবে পার্স ও প্রিন্ট করার জন্য সিনট্যাক্স সুগার সমর্থন করে। পার্সারটি -start , -update , এবং -done সাফিক্সগুলোকে বিশেষভাবে বিবেচনা করে এবং স্বয়ংক্রিয়ভাবে অ্যাসিঙ্ক্রোনাস কম্পিউটেশন ও র‍্যাপ করা ইনস্ট্রাকশনটি (সাফিক্স ছাড়া) তৈরি করে।

উদাহরণস্বরূপ, একটি অ্যাসিঙ্ক্রোনাস custom-call এভাবে লেখা যেতে পারে:

%cc-start = ((f32[64]), f32[32], s32[]) custom-call-start(%operand),
                                        custom_call_target="foo"
%cc-update = ((f32[64]), f32[32], s32[]) custom-call-update(%cc-start)
%result = f32[32] custom-call-done(%cc-update)

পার্সার এটিকে ডিসুগার করে নিম্নলিখিত সমতুল্য HLO-তে রূপান্তর করে:

%async_computation {
  %p0 = f32[64] parameter(0)
  ROOT %custom-call = f32[32] custom-call(%p0), custom_call_target="foo"
}

%async-start = ((f32[64]), f32[32], s32[]) async-start(%operand),
                calls=%async_computation
%async-update = ((f32[64]), f32[32], s32[]) async-update(%async-start)
%result = f32[32] async-done(%async-update)

এই ডিসুগারিং বেশিরভাগ HLO অপকোডের জন্য সমর্থিত (যেমন, custom-call , dot , all-reduce , ইত্যাদি)।

ব্যতিক্রম

যাতে কোনো অস্পষ্টতা তৈরি না হয়, পার্সার সেইসব অপারেশনকে ডিসুগার করবে না যেগুলোর -start এবং/অথবা -done সাফিক্স দিয়ে সুস্পষ্ট ফার্স্ট-ক্লাস অপকোড সংজ্ঞায়িত করা আছে (যেমন, copy-start / copy-done , collective-permute-start / collective-permute-done )। এগুলো তাদের নিজ নিজ ফার্স্ট-ক্লাস অপকোড ব্যবহার করতে থাকবে।

দেরিতে বাঁধাই

কিছু ক্ষেত্রে, একটি অ্যাসিঙ্ক্রোনাস অপারেশন শুরু হওয়ার সময় এর সমস্ত অপারেন্ড (ইনপুট) বা আউটপুট উপলব্ধ বা বরাদ্দ করা থাকে না। XLA লেট বাইন্ডিং সমর্থন করে, যা async-update ধাপ চলাকালীন অপারেন্ডগুলোকে পর্যায়ক্রমে বাইন্ড করতে এবং async-update বা async-done ধাপ চলাকালীন আউটপুটগুলোকে বাইন্ড করতে দেয়।

HLO-তে প্রতিনিধিত্ব

একটি আহূত গণনার জন্য যা প্রত্যাশা করে \(N\) প্যারামিটারের সংখ্যা কম হলেও, আমরা অ্যাসিঙ্ক্রোনাস এক্সিকিউশন শুরু করতে পারি। \(N\) অপারেন্ডগুলো। অবশিষ্ট অপারেন্ডগুলো পরবর্তী async-update নির্দেশাবলীতে পাঠানো হয়।

  • async-start প্রথমটিকে বাইন্ড করে। \(K\) অপারেন্ড (\(K < N\))
  • async-update নির্দেশাবলী বাকিগুলোকে সংযুক্ত করে \(N - K\) অপারেন্ড।

অপারেন্ড বাইন্ডিং অবশ্যই বাম থেকে ডানে ক্রমানুসারে হতে হবে। অর্থাৎ, যদি কোনো গণনার জন্য প্যারামিটার প্রয়োজন হয়। \(P_0, P_1, \dots, P_{N-1}\)অ্যাসিঙ্ক চেইন জুড়ে সেগুলোকে অবশ্যই সেই ক্রমেই আবদ্ধ করতে হবে।

async-start এবং async-update শেপগুলো ক্রমবর্ধমানভাবে আবদ্ধ প্যারামিটারগুলোকে প্রতিফলিত করে। বিশেষত, যত বেশি অপারেন্ড আবদ্ধ হয়, টাপল শেপের (অপারেন্ড শেপগুলোর) প্রথম উপাদানটি তত বৃদ্ধি পায়।

আউটপুট বাইন্ডিং অপারেন্ড বাইন্ডিং থেকে স্বাধীন এবং এটি অ্যাসিঙ্ক চেইনের যেকোনো ধাপে ঘটতে পারে (হয় async-update অথবা চূড়ান্ত async-done )।

kCall সহ উদাহরণ

%foo নামক একটি কল করা কম্পিউটেশন বিবেচনা করুন যা দুটি প্যারামিটার গ্রহণ করে:

%foo {
  %p0 = f32[] parameter(0)
  %p1 = f32[] parameter(1)
  ROOT %add = f32[] add(%p0, %p1)
}

আমরা এই গণনাটি অ্যাসিঙ্ক্রোনাসভাবে কল করতে পারি, শুরুতে %p0 এবং আপডেটে %p1 বাইন্ড করে:

%call-start = ((f32[]), (), s32[]) call-start(%operand0), to_apply=%foo
%call-update = ((f32[], f32[]), f32[], s32[]) call-update(%call-start, %operand1)
%result = f32[] call-done(%call-update)

পার্সার এটিকে ডিসুগার করে নিম্নলিখিত HLO-তে পরিণত করে:

%async-start = ((f32[]), (), s32[]) async-start(%operand0), calls=%foo
%async-update = ((f32[], f32[]), f32[], s32[]) async-update(%async-start, %operand1)
%result = f32[] async-done(%async-update)

বিলম্বিত-সীমাবদ্ধ আউটপুট

অপারেন্ড (ইনপুট) ছাড়াও, একটি অ্যাসিঙ্ক্রোনাস অপারেশনের আউটপুটগুলোও লেট বাইন্ড করা যেতে পারে। এটি তখন কাজে আসে যখন অপারেশন শুরুর সময় আউটপুট বাফারগুলো জানা থাকে না বা বরাদ্দ করা থাকে না।

বিলম্বিত আউটপুটগুলি উপস্থাপন করতে:

  1. async-start (বা call-start ) ইনস্ট্রাকশনটি এর আউটপুট শেপের (রেজাল্ট স্লট) ১ নম্বর ইন্ডেক্সে একটি খালি টাপল () দিয়ে সংজ্ঞায়িত করা হয়।
  2. পরবর্তী একটি async-update (বা call-update ) নির্দেশনা খালি টাপলটিকে প্রতিস্থাপন করে ইনডেক্স ১-এ প্রকৃত আউটপুট আকৃতি নির্দিষ্ট করে।
  3. বিকল্পভাবে, চেইনের শেষে async-done (বা call-done ) নির্দেশনার মাধ্যমে আউটপুটকে আবদ্ধ করা যেতে পারে, যা চূড়ান্ত আউটপুট আকৃতিটি ফেরত দেয়। চেইনে মধ্যবর্তী async-update ধাপ থাকুক বা না থাকুক, এটি করা সম্ভব।

async-update সহ উদাহরণ

// Output is not bound at start (index 1 is ())
%call-start = ((f32[1024]), (), s32[]) call-start(%input_buffer), to_apply=%foo

// Output is bound at update (index 1 becomes (f32[1024]))
%call-update = ((f32[1024]), (f32[1024]), s32[]) call-update(%call-start, %output_buffer)

%result = (f32[1024]) call-done(%call-update)

পার্সার এটিকে ডিসুগার করে নিম্নলিখিত রূপে পরিণত করে:

%async-start = ((f32[1024]), (), s32[]) async-start(%input_buffer), calls=%foo
%async-update = ((f32[1024]), (f32[1024]), s32[]) async-update(%async-start, %output_buffer)
%result = (f32[1024]) async-done(%async-update)

async-done সহ উদাহরণ ( async-update ছাড়া)

যদি কোনো মধ্যবর্তী আপডেট ধাপ না থাকে, তাহলে আউটপুটটি সরাসরি async-done এ বাইন্ড করা যেতে পারে:

// Output is not bound at start (index 1 is ())
%call-start = ((f32[1024]), (), s32[]) call-start(%input_buffer), to_apply=%foo

// Output is bound at done
%result = (f32[1024]) call-done(%call-start)

পার্সার এটিকে ডিসুগার করে নিম্নলিখিত রূপে পরিণত করে:

%async-start = ((f32[1024]), (), s32[]) async-start(%input_buffer), calls=%foo
%result = (f32[1024]) async-done(%async-start)

মধ্যবর্তী async-update এবং async-done এ আউটপুট আবদ্ধ করার উদাহরণ

যদি অপারেন্ড বাইন্ড করার জন্য মাঝখানে আপডেটের ধাপ থাকে, কিন্তু আউটপুটটি একেবারে শেষে বাইন্ড করা হয়:

// Output is not bound at start, no operands bound
%call-start = ((), (), s32[]) call-start(), to_apply=%foo

// Operands are bound at update, but output remains unbound (index 1 is ())
%call-update = ((f32[], f32[]), (), s32[]) call-update(%call-start, %operand0, %operand1)

// Output is bound at done
%result = f32[] call-done(%call-update)

পার্সার এটিকে ডিসুগার করে নিম্নলিখিত রূপে পরিণত করে:

%async-start = ((), (), s32[]) async-start(), calls=%foo
%async-update = ((f32[], f32[]), (), s32[]) async-update(%async-start, %operand0, %operand1)
%result = f32[] async-done(%async-update)