XLA:GPU অ্যাহেড-অফ-টাইম (AOT) সংকলন

অ্যাহেড-অফ-টাইম (AOT) কম্পাইলেশন একটি XLA কম্পিউটেশনকে একই পরিবেশে কম্পাইল ও সিরিয়ালাইজ করে, যার ফলে প্রাপ্ত বাইনারি আর্টিফ্যাক্টটি পরবর্তীতে আলাদাভাবে ডেপ্লয় করা একটি রানটাইম দ্বারা লোড ও এক্সিকিউট করা যায়।

জাস্ট-ইন-টাইম (JIT) কম্পাইলেশনের বিপরীতে, যা প্রথমবার চালু হওয়ার সময় চলমান প্রসেসে কম্পিউটেশনগুলো কম্পাইল করে, AOT কম্পাইলেশনকে সার্ভিং পাথ থেকে সরিয়ে দেয়, যা স্টার্টআপ ল্যাটেন্সি কমায় এবং কম্পাইলেশনকে একটি ডেডিকেটেড বিল্ড এনভায়রনমেন্টে চলতে দেয়।

ওয়ার্কফ্লো এবং এপিআই

আদর্শ AOT পাইপলাইনটি কম্পাইল, সিরিয়ালাইজ এবং লোড/ডিসেরিয়ালাইজ পর্যায়গুলো নিয়ে গঠিত:

• PjRt এন্ট্রি পয়েন্ট: C++ PjRt ইন্টারফেসে ( xla/pjrt/pjrt_client.h ), PjRtClient::Compile (বা CompileAndLoad ) এর মাধ্যমে কম্পাইলেশন শুরু করা হয়। ফলস্বরূপ এক্সিকিউটেবলটি PjRtExecutable::SerializeExecutable এর মাধ্যমে একটি স্ট্রিং-এ সিরিয়ালাইজ করা হয়। সার্ভিং হোস্টে, রানটাইম PjRtClient::DeserializeExecutable ব্যবহার করে এক্সিকিউটেবলটি পুনর্গঠন করে অথবা PjRtClient::LoadSerializedExecutable দিয়ে সরাসরি এক্সিকিউশনের জন্য লোড করে।
• কম্পাইলার-স্তরের উপাদানসমূহ: কম্পাইলার স্তরে, অ্যাহেড-অফ-টাইম কম্পাইলেশন একটি GpuAotCompilationResult (একটি CompiledModule ; AotCompilationResult একটি অপ্রচলিত বিকল্প নাম) তৈরি করে। এই অবজেক্টটি সিরিয়ালাইজড GpuExecutableProto এনক্যাপসুলেট করে এবং SerializeAsString() ও LoadExecutable() দুটিকে উন্মুক্ত করে।

• কমান্ড-লাইন টুলিং: xla_compile টুলটি HLO বা StableHLO মডিউলগুলোকে কম্পাইল করে সিরিয়ালাইজড এক্সিকিউটেবলে পরিণত করে:

  xla_compile --module_file=module.hlo --output_file=output.bin \
      --platform=gpu --gpu_target_config=gpu_target_config.pbtxt
  

  টার্গেট ডিভাইস কনফিগারেশন প্রদান করে জিপিইউ ছাড়াই হোস্টে কম্পাইলেশন চালানো যায় ( xla/backends/gpu/target_config README দেখুন)।

সিরিয়ালাইজড কী

সিরিয়ালাইজড GPU প্রোগ্রামগুলির শীর্ষ-স্তরের কন্টেইনার হল GpuExecutableProto :

• thunks : এক্সিকিউশন থাঙ্কগুলোর ক্রম। কার্নেল থাঙ্কস (উদাহরণস্বরূপ CustomKernelThunk ) তাদের নিজস্ব কম্পাইল করা কার্নেল বাইনারি (NVIDIA-এর জন্য cubin , ROCm-এর জন্য hsaco ) এমবেড করে, তাই বেশিরভাগ ডিভাইস কোড এখানেই থাকে।
• binary : একটি মডিউল-স্তরের ডিভাইস বাইনারি। বর্তমানে এটি কেবল মডিউলের ধ্রুবক বাফারগুলো ধারণ করে; এখানে আর কার্নেল সংরক্ষণ করা হয় না।
• buffer_allocations : বাফারের আকার এবং বরাদ্দের অংশ।
• gpu_compute_capability : কম্পাইল টার্গেটের হার্ডওয়্যার আর্কিটেকচার সক্ষমতা।
• executable_abi_version : টুলচেইন এবং রানটাইম ABI সংস্করণ।
• hlo_module_with_config : এমবেডেড সিরিয়ালাইজড HloModuleProto এবং ডিবাগ মেটাডেটা।

GpuExecutableProto থেকে পরোক্ষভাবে পৌঁছানো যায় এমন প্রতিটি প্রোটো মেসেজ এবং ফিল্ড , GPU AOT কম্প্যাটিবিলিটি গাইডে বর্ণিত কম্প্যাটিবিলিটি কন্ট্রাক্টের অংশ।

কীভাবে একটি আর্টিফ্যাক্ট লোড এবং এক্সিকিউট করা হয়

1. থাঙ্ক গ্রাফটি পুনর্গঠন করুন: ThunkProto এন্ট্রিগুলোকে ইন-মেমরি থাঙ্ক সিকোয়েন্সে ডিসিরিয়ালাইজ করুন এবং এই প্রক্রিয়ায় কাস্টম কল, এফএফআই, ও কার্নেল সিম্বলগুলো রিজলভ করুন। এমবেডেড এইচএলও মডিউলটিও পার্স করা হয়, তাই লোডিং বাইনারির কাছে অজানা কোনো এইচএলও অপকোড লোড প্রক্রিয়াকে ব্যর্থ করে দেয়।
2. থাঙ্ক গ্রাফ রূপান্তর ধাপসমূহ চালান: ইন-মেমরি থাঙ্কগুলোর উপর লোড-টাইম অপ্টিমাইজেশন এবং রূপান্তর ধাপসমূহ সম্পাদন করুন (উদাহরণস্বরূপ, CUDA গ্রাফের মাধ্যমে কাজ জমা দেওয়ার জন্য রানটাইম CommandBufferThunk তৈরি করা)। CommandBufferThunk লোড হওয়ার সময় ডায়নামিকভাবে পুনর্গঠিত হয় এবং এটি কখনোই সরাসরি সিরিয়ালাইজ করা হয় না।
3. বাফার বরাদ্দ করুন: ক্রমিক বাফার অ্যাসাইনমেন্ট স্লাইস অনুযায়ী মেমরি স্পেস এবং বাফার বরাদ্দ শুরু করুন।
4. কার্যকর করুন: ডিভাইস স্ট্রিমগুলিতে প্রস্তুত থাঙ্ক সিকোয়েন্সটি প্রেরণ করুন।

পাস, ফিউশন, কোডজেন এবং অটোটিউনিং-এর এমন পরিবর্তনসমূহ, যা সিরিয়ালাইজড স্কিমা ও রানটাইম সিম্যান্টিকস অক্ষুণ্ণ রাখে, তা AOT কম্প্যাটিবিলিটিকে প্রভাবিত করে না।

সমস্যা সমাধান: লোড ব্যর্থতা যা ভার্সনিং বাগ নয়

সব এক্সিকিউটেবল লোড ব্যর্থতার কারণ সিরিয়ালাইজেশন রিগ্রেশন নয়। দুটি সাধারণ পরিবেশগত অমিল লোড-টাইম ত্রুটি তৈরি করে:

• কম্পিউট ক্যাপাবিলিটি অমিল: GpuExecutable::FromProto যাচাই করে যে সিরিয়ালাইজড gpu_compute_capability টার্গেট ডিভাইসের সাথে মেলে কিনা। NVIDIA Hopper-এর জন্য কম্পাইল করা একটি এক্সিকিউটেবল NVIDIA Ampere-এ চলতে পারে না।
• CUDA টুলচেইন ABI অমিল: CudaRuntimeAbiVersion::IsCompatibleWith পরীক্ষা করে দেখে যে রানটাইম হোস্টের CUDA টুলকিট, cuDNN, এবং CUB সংস্করণগুলি ExecutableAbiVersionProto তে নথিভুক্ত ন্যূনতম সংস্করণগুলিকে পূরণ করে বা অতিক্রম করে কিনা। যদি হোস্ট ড্রাইভার বা লাইব্রেরি পরিবেশ কম্পাইল টার্গেটের চেয়ে পুরোনো হয়, তাহলে FailedPreconditionError সাথে লোডিং ব্যর্থ হয়।