HLO থেকে Thunks

প্রাক-অপ্টিমাইজেশান HLO

আমরা প্রাক-অপ্টিমাইজেশন HLO দিয়ে শুরু করি। প্রাক-অপ্টিমাইজেশান এইচএলও-তে এমন অপ্স নেই যা XLA-এর অভ্যন্তরীণ বলে বিবেচিত হয়, যেমন fusion বা bitcast । Ops-এর এই পর্যায়ে কোনো লেআউট নেই, অথবা যদি থাকে, তাহলে তা উপেক্ষা করা হবে। প্রাক-অপ্টিমাইজেশান HLO সাধারণত Tensorflow এবং JAX এর মত উচ্চ স্তরের ফ্রেমওয়ার্ক দ্বারা উত্পাদিত হয়। XLA পতাকা -xla_dump_to ব্যবহার করার সময়, প্রাক-অপ্টিমাইজেশান HLO ফাইলের নাম প্রত্যয় "before_optimizations.txt" সহ একটি ফাইলে ডাম্প করা হয়।

HLO মডিউল অপ্টিমাইজ করুন

XLA:GPU পাইপলাইন পাসের একটি ক্রম চালিয়ে প্রাক-অপ্টিমাইজেশন HLO কে অপ্টিমাইজ করা HLO-তে পরিণত করবে। পাসগুলি শব্দার্থগতভাবে একসাথে গোষ্ঠীবদ্ধ করা যেতে পারে এবং নিম্নলিখিত ক্রমে চালানো যেতে পারে:

শেয়ারিং সম্পর্কিত পাস

শার্ডি পার্টিশনার বা SPMD শার্ডিং।

অপ্টিমাইজেশান পাস।

এতে বৈধকরণ পাস এবং সরলীকরণ পাস উভয়ই অন্তর্ভুক্ত থাকতে পারে।

সমষ্টিগত অপ্টিমাইজেশান পাস.

অপ্টিমাইজেশান পাসের অনুরূপ, কিন্তু যৌথ অপারেশনগুলিতে ফোকাস করে৷

লেআউট অ্যাসাইনমেন্ট পাস

প্রতিটি HLO অপকে একটি লেআউট বরাদ্দ করা হয় যা নির্দেশের আকারের একটি অংশ। বিন্যাস নিয়ন্ত্রণ করে কিভাবে টেনসর মেমরিতে শারীরিকভাবে বিছানো হয়।

লেআউট সহ একটি আকৃতির উদাহরণ:

f32[10,20,30]{2,0,1}

উপাদানের প্রকারের পরে, আকৃতির যৌক্তিক মাত্রা রয়েছে, তারপরে ছোট থেকে বড় ক্রমে বিন্যাস স্থানান্তর। এই উদাহরণে, সবচেয়ে ছোট মাত্রা হল 30, দ্বিতীয় সবচেয়ে ছোট মাত্রা হল 10, এবং প্রধান মাত্রা হল 20৷

লেআউট অ্যাসাইনমেন্টের লক্ষ্য হল একটি লোভী কৌশল ব্যবহার করে প্রয়োজনীয় শারীরিক স্থানান্তরের সংখ্যা হ্রাস করা। এটি নির্দিষ্ট লেআউট সীমাবদ্ধতার সাথে শুরু হয় (যেমন CuDNN/cuBLAS লাইব্রেরিগুলি পরপর মাত্রা আশা করে) এবং লেআউটকে "নিচে" প্রচার করে এবং তারপরে HLO গ্রাফটিকে "উপরে" করে। লেআউট প্রচারের শেষে, কিছু নির্দেশে বিরোধপূর্ণ লেআউট থাকতে পারে, একটি অপারেন্ড থেকে প্রচারিত, একটি ব্যবহারকারীর কাছ থেকে প্রচারিত। এই বিরোধের সমাধান করার জন্য, একটি copy HLO নির্দেশ সন্নিবেশ করা হয় যা অপারেন্ড লেআউট থেকে নির্দেশ লেআউটে লেআউট পরিবর্তন করে।

বিন্যাস স্বাভাবিককরণ পাস

প্রদত্ত যে শারীরিক আকৃতি বের করা কিছুটা কঠিন, লেআউট স্বাভাবিককরণ আকৃতিটিকে পুনরায় লেখার চেষ্টা করে যাতে এটি ডিফল্ট লেআউট ব্যবহার করে {rank-1, rank-2, …, 0} । উপরের উদাহরণে, স্বাভাবিক আকার হবে f32[20,10,30]{2,1,0} । transpose + bitcast সংমিশ্রণে লেআউট পরিবর্তন করে এমন অপস অনুলিপি করা হয়। প্রদত্ত যে বর্তমানে আমরা সমস্ত অপ্সকে স্বাভাবিক করতে পারি না, এখনও কিছু অপ্স আছে যেগুলিতে নন-ডিফল্ট লেআউট থাকতে পারে, বিশেষ করে gather এবং dot । স্বাভাবিকীকৃত অপস এবং নন-নর্মালাইজড অপস এর মধ্যে সীমারেখায় bitcast অপস থাকবে যা একটি ট্রান্সপোজকে প্রতিনিধিত্ব করে, অর্থাৎ একটি লেআউটের সাথে একটি ট্রান্সপোজ যা এটিকে শারীরিকভাবে একটি নো-অপ করে তোলে।

লেআউট স্বাভাবিকীকরণ কিছু অন্তর্নিহিত ট্রান্সপোজকেও সুস্পষ্ট করে তোলে যা গুরুত্বপূর্ণ কারণ কোডজেন একটি ডেডিকেটেড ইমিটার দিয়ে স্পষ্ট ট্রান্সপোজ পরিচালনা করতে পারে। উদাহরণস্বরূপ, একটি পুনঃআকৃতি প্রযুক্তিগতভাবে অপারেন্ড এবং ফলাফলের মধ্যে একটি ভিন্ন শারীরিক বিন্যাস থাকার অনুমতি দেওয়া হয় (যেমন ভিন্ন র্যাঙ্কের কারণে)। ReshapeDecomposer পাস যেটি লেআউট স্বাভাবিকীকরণ পাসের অংশ হিসাবে চালিত হয় তা একটি রিশেপকে transpose , রিশেপ bitcast এবং transpose একটি সিকোয়েন্সে পরিণত করে।

পোস্ট লেআউট অ্যাসাইনমেন্ট অপ্টিমাইজেশান পাস

এখানে সবচেয়ে গুরুত্বপূর্ণ পাস হল Triton fusions (GEMM fusions + Softmax/layernorm fusions) অথবা লাইব্রেরী কলে পুনরায় লেখা। কিন্তু অটোটিউনিংও এই ধাপে চলে, যেখানে আমরা কনভোলিউশন বা বিন্দুর জন্য সেরা অ্যালগরিদম বা লিগ্যাসি ট্রাইটন জিইএমএম ইমিটার দ্বারা পরিচালিত ডটগুলির জন্য সেরা টাইলিং বা একটি নির্দিষ্ট ডট ফিউশনের জন্য আমাদের ট্রাইটন বা কিউব্লাস ব্যবহার করা উচিত কিনা তা বেছে নিই।

ফিউশন পাস

দুটি প্রধান পাস হল PriorityFusion এবং Multi-Output ফিউশন।

PriorityFusion এ, আমরা খরচ মডেল দ্বারা পরিচালিত ফিউশন গঠন করি। ফিউজ করার সময় আমরা বেশ কয়েকটি ব্যবহারকারীর সাথে ডুপ্লিকেট অপের অনুমতি দেব যদি অপটি সমস্ত ব্যবহারকারীর মধ্যে মিশ্রিত করা যায়। আমরা যদি সম্ভব হয় তবে বিদ্যমান ট্রাইটন সফটম্যাক্স ফিউশনগুলিকে প্রসারিত করার অনুমতি দেব।

Multi-Output ফিউশন হল একটি পৃথক পাস যা অপস/ফিউশনকে একসাথে ফিউজ করার অনুমতি দেয় যা একটি অপারেন্ড ভাগ করে, অথবা ফিউজ অপারেন্ড/অপারেন্ড ফিউশনকে ডুপ্লিকেশন ছাড়াই ব্যবহারকারীদের মধ্যে অতিরিক্ত আউটপুট(গুলি) যোগ করে যাতে অপের অন্যান্য ব্যবহারকারীদের এই আউটপুটে রিডাইরেক্ট করা যায়। এই পাসটি HLO গ্রাফে চক্র প্রবর্তন না করার জন্য সতর্কতা অবলম্বন করা প্রয়োজন।

মাল্টি-আউটপুট ফিউশনের পরে, আমরা কমন সাবএক্সপ্রেশন এলিমিনেশন চালাই ( HloCSE পাস) যা একই ফিউশনে শেষ হলে পূর্বে ডুপ্লিকেট করা অপ্সগুলিকে আবার একসাথে একত্রিত করবে।

বেশ কিছু পোস্ট-ফিউশন পাস

সমষ্টির সাথে সম্পর্কিত বেশ কয়েকটি পাস (যেমন এগুলিকে অ্যাসিঙ্কে পরিণত করা, বা সমষ্টির একটি নির্দিষ্ট আপেক্ষিক আদেশ প্রয়োগ করা)।

অবশেষে আমরা CopyInsertion চালাই যেখানে কপিগুলি যোগ করা হয় তা নিশ্চিত করার জন্য যে ইন-প্লেস ক্রিয়াকলাপগুলি এখনও অন্য কোথাও প্রয়োজনীয় ডেটা ওভাররাইট করে না।

অপ্টিমাইজেশনের শেষে, অপ্টিমাইজ করা HLO ডাম্প করা হয় যদি ফ্ল্যাগ -xla_dump_to ব্যবহার করে এমন একটি ফাইলে যার ফাইলের নাম প্রত্যয় "after_optimizations.txt" থাকে। আপনি যদি মধ্যবর্তী পাসের পরে HLO ডাম্প করতে চান যা আসলে HloModule পরিবর্তন করে, আপনি ফ্ল্যাগ -xla_dump_hlo_pass_re=.* (বা নির্দিষ্ট পাসে সীমাবদ্ধ করার জন্য একটি নির্দিষ্ট রেগুলার এক্সপ্রেশন) ব্যবহার করতে পারেন।

সময়সূচী

সময়সূচী ছাড়া একটি HloModule এখনও কিছু মাত্রার স্বাধীনতা রয়েছে যাতে অপারেশনগুলি প্রক্রিয়া করা হয়। মূলত অপারেন্ড/ফলাফল সম্পর্ক এবং নিয়ন্ত্রণ নির্ভরতা অনুযায়ী যেকোনো টপোলজিক্যাল সাজানো ঠিক আছে। সময়সূচী একটি নির্দিষ্ট আদেশ কার্যকর করে। এটি প্রয়োজনীয় মেমরির পরিমাণকে প্রভাবিত করে, কারণ আমরা একটি বাফার পুনঃব্যবহার করতে পারি না যতক্ষণ না সেই বাফারের সমস্ত পাঠক প্রক্রিয়া করা না হয়। প্রাথমিক ধাপে, আমরা বিভিন্ন সময়সূচী অ্যালগরিদম চেষ্টা করি এবং সময়সূচী বাছাই করি যা পিক মেমরি খরচ কমিয়ে দেয়।

ফলো-আপ হিসাবে, আমরা LatencyHidingScheduler পাসটি চালাই যা কম্পিউট-কমিউনিকেশন ওভারল্যাপকে সর্বাধিক করার চেষ্টা করে কিন্তু মেমরির ব্যবহার আবার বাড়িয়ে দিতে পারে।

শিডিউল করার পরে, আমরা HloRematerialization চালাই যা মেমরির ব্যবহার কমানোর চেষ্টা করে যদি পিক মেমরি খরচ আমাদের উপলব্ধ মেমরির পরিমাণের চেয়ে বেশি হয়। এটি পারফরম্যান্সের খরচে, যেমন কিছু ফিউশন বিভক্ত হতে পারে এবং কিছু অপ্স ছোট বাফার লাইফটাইম থাকার জন্য ডুপ্লিকেট হতে পারে। যদি রিমেটেরিয়ালাইজেশন ঘটতে থাকে, তাহলে মডেল সাইডে প্রয়োজনীয় মেমরির পরিমাণ কমানোর উপায় আছে কিনা তা দেখা সম্ভবত অর্থপূর্ণ হবে (যেমন ছোট ব্যাচের আকার)।

Thunks এবং CommandBuffers

টিবিডি

বাফার অ্যাসাইনমেন্ট

আমরা LLVM IR-এ নামানোর ঠিক আগে, আমরা বাফার অ্যাসাইনমেন্ট পাসগুলি চালাই যা HLO গ্রাফের প্রতিটি নির্দেশে বাফার স্লাইস বরাদ্দ করবে। বাফার অ্যাসাইনমেন্ট বিভিন্ন ধাপে সঞ্চালিত হয়।

1. HloDataflowAnalysis নির্দেশাবলীতে HloValues ​​(মূলত লজিক্যাল বাফার) বরাদ্দ করে। ইন-প্লেস অপারেশনের জন্য, একটি অপারেন্ডের HloValue পুনরায় ব্যবহার করা যেতে পারে। একটি বিকল্প একাধিক HloValue সংজ্ঞায়িত করতে পারে (যেমন একটি tuple ফলাফল আকৃতি সহ)।

2. HloAliasAnalysis অ্যালিয়াসিং ক্রিয়াকলাপগুলির জন্য বাফারগুলিকে একত্রিত করার চেষ্টা করে এবং HloValue থেকে HloBuffer এ একটি ম্যাপিং গণনা করে।

3. BufferAssignment HloBuffers এর একটি ম্যাপিং একটি বড় বাফারের ভিতরে বাফার স্লাইসকে এমনভাবে গণনা করে যাতে একই বাফার স্লাইস ওভারল্যাপিং লাইফ টাইম সহ বিভিন্ন HloBuffers জন্য ব্যবহার করা হয় না। অপ্সের জন্য যেগুলি উপনাম হতে পারে, এটি ঠিক আছে যে একটি সামান্য ওভারল্যাপ আছে (একটি HloBuffer এর শেষ সময় অন্য HloBuffer এর শুরুর সময়ের সাথে মিলে যেতে পারে)৷ পতাকা -xla_dump_to ব্যবহার করার সময়, বাফার অ্যাসাইনমেন্ট সম্পর্কে কিছু তথ্য "after_optimizations-buffer-assignment.txt" নামের প্রত্যয় সহ একটি ফাইলে ডাম্প করা হয়।