XLA:GPU-তে HLO-এর জন্য কোড তৈরি করার তিনটি উপায় রয়েছে।
- বাহ্যিক লাইব্রেরিতে কাস্টম কল দিয়ে HLO প্রতিস্থাপন করা হচ্ছে, যেমন NVidia cuBLAS , cuDNN ।
- ব্লক-লেভেলে HLO টাইলিং করুন এবং তারপর OpenAI Triton ব্যবহার করুন।
- XLA ইমিটার ব্যবহার করে ক্রমশ HLO-কে LLVM IR থেকে কমানো।
এই নথিটি XLA:GPU ইমিটারের উপর দৃষ্টি নিবদ্ধ করা হয়েছে।
নায়ক-ভিত্তিক কোডজেন
XLA:GPU-তে 7টি ইমিটার প্রকার রয়েছে। প্রতিটি ইমিটার টাইপ ফিউশনের একটি "হিরো" এর সাথে মিলে যায়, অর্থাৎ ফিউজড কম্পিউটেশনের সবচেয়ে গুরুত্বপূর্ণ অপশন যা পুরো ফিউশনের জন্য কোড জেনারেশনকে আকার দেয়।
উদাহরণস্বরূপ, ট্রানপোজ ইমিটার নির্বাচন করা হবে যদি ফিউশনের মধ্যে একটি HloTransposeInstruction থাকে যার জন্য মেমরি পড়ার এবং লেখার ধরণগুলিকে উন্নত করতে শেয়ার্ড মেমরি ব্যবহার করা প্রয়োজন। রিডাকশন ইমিটার শাফেল এবং শেয়ার করা মেমরি ব্যবহার করে রিডাকশন জেনারেট করে। লুপ ইমিটার হল ডিফল্ট ইমিটার। যদি একটি ফিউশনে একটি হিরো না থাকে যার জন্য আমাদের একটি বিশেষ ইমিটার আছে, তাহলে লুপ ইমিটার ব্যবহার করা হবে।
উচ্চ-স্তরের ওভারভিউ
কোডটি নিম্নলিখিত বড় বিল্ডিং ব্লকগুলি নিয়ে গঠিত:
- কম্পিউটেশন পার্টিশনার - একটি HLO ফিউশন কম্পিউটেশনকে ফাংশনে বিভক্ত করা
- ইমিটারস - পার্টিশন করা এইচএলও ফিউশনকে এমএলআইআর (
xla_gpu,tensor,arith,math,scfউপভাষায়) রূপান্তর করা - কম্পাইলেশন পাইপলাইন - অপ্টিমাইজ করে এবং IR কে LLVM এ কম করে
বিভাজন
computation_partitioner.h দেখুন।
উপাদানবিহীন এইচএলও নির্দেশ সবসময় একসাথে নির্গত করা যায় না। নিম্নলিখিত HLO গ্রাফ বিবেচনা করুন:
param
|
log
| \
| transpose
| /
add
যদি আমরা এটি একটি একক ফাংশনে নির্গত করি, log add প্রতিটি উপাদানের জন্য দুটি ভিন্ন সূচকে অ্যাক্সেস করা হবে। পুরানো নির্গমনকারীরা দুইবার log তৈরি করে এই সমস্যার সমাধান করে। এই নির্দিষ্ট গ্রাফের জন্য, এটি একটি সমস্যা নয়, কিন্তু যখন একাধিক বিভাজন থাকে, তখন কোডের আকার দ্রুতগতিতে বৃদ্ধি পায়।
এখানে, আমরা গ্রাফটিকে টুকরো টুকরো করে বিভাজন করে এই সমস্যার সমাধান করি যা নিরাপদে একটি ফাংশন হিসাবে নির্গত হতে পারে। মানদণ্ড হল:
- শুধুমাত্র একজন ব্যবহারকারীর নির্দেশাবলী তাদের ব্যবহারকারীর সাথে একসাথে নির্গত করা নিরাপদ।
- একাধিক ব্যবহারকারীর নির্দেশাবলী তাদের ব্যবহারকারীদের সাথে একসাথে নির্গত করা নিরাপদ যদি সেগুলি সমস্ত ব্যবহারকারীর দ্বারা একই সূচকের মাধ্যমে অ্যাক্সেস করা হয়।
উপরের উদাহরণে, log বিভিন্ন সূচক add এবং tranpose অ্যাক্সেস করে, তাই তাদের সাথে এটি নির্গত করা নিরাপদ নয়।
তাই গ্রাফটিকে তিনটি ফাংশনে বিভক্ত করা হয়েছে (প্রতিটিতে মাত্র একটি নির্দেশ রয়েছে)।
slice এবং add এর pad সহ নিম্নলিখিত উদাহরণের ক্ষেত্রেও এটি প্রযোজ্য।
মৌলিক নির্গমন
মৌলিক_hlo_to_mlir.h দেখুন।
মৌলিক নির্গমন HloInstructions এর জন্য লুপ এবং গণিত/অ্যারিথ অপ্স তৈরি করে। বেশিরভাগ অংশের জন্য, এটি সোজা, কিন্তু এখানে কিছু আকর্ষণীয় জিনিস চলছে।
ইনডেক্সিং রূপান্তর
কিছু নির্দেশনা ( transpose , broadcast , reshape , slice , reverse এবং আরও কিছু ) হল সূচকে বিশুদ্ধভাবে রূপান্তর: ফলাফলের একটি উপাদান তৈরি করতে, আমাদের ইনপুটের অন্য কিছু উপাদান তৈরি করতে হবে। এর জন্য, আমরা XLA এর indexing_analysis পুনঃব্যবহার করতে পারি, যেটিতে একটি নির্দেশের জন্য ইনপুট ম্যাপিং থেকে আউটপুট তৈরি করার ফাংশন রয়েছে।
উদাহরণস্বরূপ, [20,40] থেকে [40,20] থেকে একটি transpose জন্য, এটি নিম্নলিখিত সূচীকরণ মানচিত্র তৈরি করবে (ইনপুট মাত্রা প্রতি একটি অ্যাফাইন এক্সপ্রেশন; d0 এবং d1 হল আউটপুট মাত্রা):
(d0, d1) -> d1
(d0, d1) -> d0
সুতরাং এই বিশুদ্ধ সূচক রূপান্তর নির্দেশাবলীর জন্য, আমরা কেবল মানচিত্রটি পেতে পারি, এটি আউটপুট সূচকগুলিতে প্রয়োগ করতে পারি এবং ফলাফল সূচকে ইনপুট তৈরি করতে পারি।
একইভাবে, pad অপটি বেশিরভাগ বাস্তবায়নের জন্য ইন্ডেক্সিং মানচিত্র এবং সীমাবদ্ধতা ব্যবহার করে। আমরা ইনপুট বা প্যাডিং মানের একটি উপাদান ফেরত দিই কিনা তা দেখার জন্য কিছু অতিরিক্ত চেক সহ pad একটি সূচীকরণ রূপান্তর।
টিপলস
আমরা অভ্যন্তরীণ tuple এস সমর্থন করি না। আমরা নেস্টেড টিপল আউটপুট সমর্থন করি না। এই বৈশিষ্ট্যগুলি ব্যবহার করে এমন সমস্ত XLA গ্রাফগুলি এমন গ্রাফগুলিতে রূপান্তরিত হতে পারে যা করে না৷
জড়ো করা
gather_simplifier দ্বারা উত্পাদিত হিসাবে আমরা শুধুমাত্র ক্যানোনিকাল সংগ্রহ সমর্থন করি।
সাবগ্রাফ ফাংশন
%p0 থেকে %p_n পরামিতি সহ একটি গণনার একটি সাবগ্রাফের জন্য এবং র্যাঙ্ক r এবং উপাদানের ধরন সহ সাবগ্রাফ রুট ( e0 থেকে e_m ), আমরা নিম্নলিখিত MLIR ফাংশন স্বাক্ষর ব্যবহার করি:
(%p0: tensor<...>, %p1: tensor<...>, ..., %pn: tensor<...>,
%i0: index, %i1: index, ..., %i_r-1: index) -> (e0, ..., e_m)
অর্থাৎ, প্রতি গণনা প্যারামিটারে আমাদের একটি টেনসর ইনপুট, আউটপুটের মাত্রা প্রতি একটি সূচক ইনপুট এবং প্রতি আউটপুটে একটি ফলাফল রয়েছে।
একটি ফাংশন নির্গত করার জন্য, আমরা কেবল উপরের এলিমেন্টাল ইমিটার ব্যবহার করি এবং সাবগ্রাফের প্রান্তে না পৌঁছানো পর্যন্ত এর অপারেন্ডগুলি পুনরাবৃত্তিমূলকভাবে নির্গত করি। তারপর, আমরা:পরামিটারের জন্য একটি tensor.extract নির্গত করি বা অন্যান্য সাবগ্রাফের জন্য একটি func.call নির্গত করি
এন্ট্রি ফাংশন
প্রতিটি বিকিরণকারীর ধরন এটি কীভাবে এন্ট্রি ফাংশন তৈরি করে, অর্থাৎ নায়কের জন্য ফাংশন আলাদা করে। এন্ট্রি ফাংশন উপরের ফাংশন থেকে ভিন্ন, যেহেতু এটিতে ইনপুট হিসাবে কোন সূচক নেই (শুধু থ্রেড এবং ব্লক আইডি) এবং আসলে কোথাও আউটপুট লিখতে হবে। লুপ ইমিটারের জন্য, এটি মোটামুটি সোজা, কিন্তু ট্রান্সপোজ এবং রিডাকশন ইমিটারগুলির অ-তুচ্ছ লেখার যুক্তি থাকে।
এন্ট্রি গণনার স্বাক্ষর হল:
(%p0: tensor<...>, ..., %pn: tensor<...>,
%r0: tensor<...>, ..., %rn: tensor<...>) -> (tensor<...>, ..., tensor<...>)
যেখানে আগের মত, %pn s হল গণনার পরামিতি, এবং %rn s হল গণনার ফলাফল। এন্ট্রি গণনা ফলাফলগুলিকে টেনসর হিসাবে নেয়, tensor.insert s আপডেটগুলি তাদের মধ্যে দেয় এবং তারপরে সেগুলি ফেরত দেয়৷ আউটপুট টেনসরের অন্য কোন ব্যবহার অনুমোদিত নয়।
সংকলন পাইপলাইন
লুপ ইমিটার
loop_mlir.h দেখুন।
আসুন GELU ফাংশনের জন্য HLO ব্যবহার করে MLIR সংকলন পাইপলাইনের সবচেয়ে গুরুত্বপূর্ণ পাসগুলি অধ্যয়ন করি।
এই HLO কম্পিউটেশনে শুধুমাত্র elementwise ops, constants এবং সম্প্রচার আছে। এটি লুপ ইমিটার ব্যবহার করে নির্গত হবে।
MLIR রূপান্তর
MLIR-এ রূপান্তর করার পরে আমরা একটি xla_gpu.loop পাই যা %thread_id_x এবং %block_id_x উপর নির্ভর করে এবং লুপটি সংজ্ঞায়িত করে যা আউটপুটের সমস্ত উপাদানকে একত্রিত লেখার গ্যারান্টি দিতে রৈখিকভাবে অতিক্রম করে।
এই লুপ প্রতিটি পুনরাবৃত্তির উপর আমরা কল
%pure_call = xla_gpu.pure_call @gelu(%input, %dim0, %dim1, %dim2)
: (tensor<6x512x4096xbf16>, index, index, index) -> bf16
রুট অপারেশন উপাদান গণনা. দ্রষ্টব্য, আমাদের কাছে @gelu এর জন্য শুধুমাত্র একটি রূপরেখাযুক্ত ফাংশন রয়েছে, কারণ বিভাজনকারী একটি টেনসর সনাক্ত করেনি যাতে 2 বা তার বেশি বিভিন্ন অ্যাক্সেস প্যাটার্ন রয়েছে।
#map = #xla_gpu.indexing_map<"(th_x, bl_x)[vector_index] -> ("
"bl_x floordiv 4096, (bl_x floordiv 8) mod 512, (bl_x mod 8) * 512 + th_x * 4 + vector_index),"
"domain: th_x in [0, 127], bl_x in [0, 24575], vector_index in [0, 3]">
func.func @main(%input: tensor<6x512x4096xbf16> , %output: tensor<6x512x4096xbf16>)
-> tensor<6x512x4096xbf16> {
%thread_id_x = gpu.thread_id x {xla.range = [0 : index, 127 : index]}
%block_id_x = gpu.block_id x {xla.range = [0 : index, 24575 : index]}
%xla_loop = xla_gpu.loop (%thread_id_x, %block_id_x)[%vector_index] -> (%dim0, %dim1, %dim2)
in #map iter_args(%iter = %output) -> (tensor<6x512x4096xbf16>) {
%pure_call = xla_gpu.pure_call @gelu(%input, %dim0, %dim1, %dim2)
: (tensor<6x512x4096xbf16>, index, index, index) -> bf16
%inserted = tensor.insert %pure_call into %iter[%dim0, %dim1, %dim2] : tensor<6x512x4096xbf16>
xla_gpu.yield %inserted : tensor<6x512x4096xbf16>
}
return %xla_loop : tensor<6x512x4096xbf16>
}
func.func private @gelu(%arg0: tensor<6x512x4096xbf16>, %i: index, %j: index, %k: index) -> bf16 {
%cst = arith.constant 5.000000e-01 : bf16
%cst_0 = arith.constant 1.000000e+00 : bf16
%cst_1 = arith.constant 7.968750e-01 : bf16
%cst_2 = arith.constant 4.467770e-02 : bf16
%extracted = tensor.extract %arg0[%i, %j, %k] : tensor<6x512x4096xbf16>
%0 = arith.mulf %extracted, %extracted : bf16
%1 = arith.mulf %0, %extracted : bf16
%2 = arith.mulf %1, %cst_2 : bf16
%3 = arith.addf %extracted, %2 : bf16
%4 = arith.mulf %3, %cst_1 : bf16
%5 = math.tanh %4 : bf16
%6 = arith.addf %5, %cst_0 : bf16
%7 = arith.mulf %6, %cst : bf16
%8 = arith.mulf %extracted, %7 : bf16
return %8 : bf16
}
ইনলাইনার
@gelu ইনলাইন করার পরে, আমরা একটি একক @main ফাংশন পাই। এটি ঘটতে পারে যে একই ফাংশনকে দুইবার বা তার বেশি বলা হয়। এই ক্ষেত্রে আমরা ইনলাইন না. ইনলাইনিং নিয়ম সম্পর্কে আরো বিস্তারিত জানা যাবে xla_gpu_dialect.cc এ।
func.func @main(%arg0: tensor<6x512x4096xbf16>, %arg1: tensor<6x512x4096xbf16>) -> tensor<6x512x4096xbf16> {
...
%thread_id_x = gpu.thread_id x {xla.range = [0 : index, 127 : index]}
%block_id_x = gpu.block_id x {xla.range = [0 : index, 24575 : index]}
%xla_loop = xla_gpu.loop (%thread_id_x, %block_id_x)[%vector_index] -> (%dim0, %dim1, %dim2)
in #map iter_args(%iter = %output) -> (tensor<6x512x4096xbf16>) {
%extracted = tensor.extract %input[%dim0, %dim1, %dim2] : tensor<6x512x4096xbf16>
%0 = arith.mulf %extracted, %extracted : bf16
%1 = arith.mulf %0, %extracted : bf16
%2 = arith.mulf %1, %cst : bf16
%3 = arith.addf %extracted, %2 : bf16
%4 = arith.mulf %3, %cst_0 : bf16
%5 = math.tanh %4 : bf16
%6 = arith.addf %5, %cst_1 : bf16
%7 = arith.mulf %6, %cst_2 : bf16
%8 = arith.mulf %extracted, %7 : bf16
%inserted = tensor.insert %8 into %iter[%dim0, %dim1, %dim2] : tensor<6x512x4096xbf16>
xla_gpu.yield %inserted : tensor<6x512x4096xbf16>
}
return %xla_loop : tensor<6x512x4096xbf16>
}
xla_gpu থেকে scf রূপান্তর
Lower_xla_gpu_to_scf.cc দেখুন।
xla_gpu.loop ভিতরে একটি সীমানা চেক সহ একটি লুপ নেস্ট উপস্থাপন করে। যদি লুপ ইনডাকশন ভেরিয়েবলগুলি ইন্ডেক্সিং ম্যাপ ডোমেনের সীমার বাইরে থাকে, তাহলে এই পুনরাবৃত্তিটি এড়িয়ে যাবে। এর মানে হল, লুপটিকে 1 বা তার বেশি নেস্টেড scf.for ops-এ একটি scf.if ভিতরে রূপান্তর করা হয়েছে।
%xla_loop = scf.for %vector_index = %c0 to %c4 step %c1 iter_args(%iter = %output) -> (tensor<6x512x4096xbf16>) {
%2 = arith.cmpi sge, %thread_id_x, %c0 : index
%3 = arith.cmpi sle, %thread_id_x, %c127 : index
%4 = arith.andi %2, %3 : i1
%5 = arith.cmpi sge, %block_id_x, %c0 : index
%6 = arith.cmpi sle, %block_id_x, %c24575 : index
%7 = arith.andi %5, %6 : i1
%inbounds = arith.andi %4, %7 : i1
%9 = scf.if %inbounds -> (tensor<6x512x4096xbf16>) {
%dim0 = xla_gpu.apply_indexing #map(%thread_id_x, %block_id_x)[%vector_index]
%dim1 = xla_gpu.apply_indexing #map1(%thread_id_x, %block_id_x)[%vector_index]
%dim2 = xla_gpu.apply_indexing #map2(%thread_id_x, %block_id_x)[%vector_index]
%extracted = tensor.extract %input[%dim0, %dim1, %dim2] : tensor<6x512x4096xbf16>
// ... more arithmetic operations
%29 = arith.mulf %extracted, %28 : bf16
%inserted = tensor.insert %29 into %iter[%dim0, %dim1, %dim2] : tensor<6x512x4096xbf16>
scf.yield %inserted : tensor<6x512x4096xbf16>
} else {
scf.yield %iter : tensor<6x512x4096xbf16>
}
scf.yield %9 : tensor<6x512x4096xbf16>
}
সমতল টেনসর
flatten_tensors.cc দেখুন।
Nd tensors 1D সম্মুখে অভিক্ষিপ্ত হয়. এটি ভেক্টরাইজেশন এবং LLVM-এ কম করাকে সহজ করবে কারণ প্রতিটি টেনসর অ্যাক্সেস এখন মেমরিতে ডেটা কীভাবে সারিবদ্ধ করা হয় তার সাথে মিলে যায়।
#map = #xla_gpu.indexing_map<"(th_x, bl_x, vector_index) -> (th_x * 4 + bl_x * 512 + vector_index),"
"domain: th_x in [0, 127], bl_x in [0, 24575], vector_index in [0, 3]">
func.func @main(%input: tensor<12582912xbf16>, %output: tensor<12582912xbf16>) -> tensor<12582912xbf16> {
%xla_loop = scf.for %vector_index = %c0 to %c4 step %c1 iter_args(%iter = %output) -> (tensor<12582912xbf16>) {
%dim = xla_gpu.apply_indexing #map(%thread_id_x, %block_id_x, %vector_index)
%extracted = tensor.extract %input[%dim] : tensor<12582912xbf16>
%2 = arith.mulf %extracted, %extracted : bf16
%3 = arith.mulf %2, %extracted : bf16
%4 = arith.mulf %3, %cst_2 : bf16
%5 = arith.addf %extracted, %4 : bf16
%6 = arith.mulf %5, %cst_1 : bf16
%7 = math.tanh %6 : bf16
%8 = arith.addf %7, %cst_0 : bf16
%9 = arith.mulf %8, %cst : bf16
%10 = arith.mulf %extracted, %9 : bf16
%inserted = tensor.insert %10 into %iter[%dim] : tensor<12582912xbf16>
scf.yield %inserted : tensor<12582912xbf16>
}
return %xla_loop : tensor<12582912xbf16>
}
ভেক্টরাইজেশন
vectorize_loads_stores.cc দেখুন।
পাসটি tensor.extract এবং tensor.insert ops-এ সূচকগুলি বিশ্লেষণ করে এবং যদি সেগুলি xla_gpu.apply_indexing দ্বারা উত্পাদিত হয় যা উপাদানগুলিকে %vector_index এ সংলগ্নভাবে অ্যাক্সেস করে এবং অ্যাক্সেসটি সারিবদ্ধ করা হয়, তাহলে tensor.extract vector.transfer_read এ রূপান্তরিত হয় এবং লুপ আউট উত্তোলন.
এই বিশেষ ক্ষেত্রে, একটি ইন্ডেক্সিং মানচিত্র রয়েছে (th_x, bl_x, vector_index) -> (th_x * 4 + bl_x * 512 + vector_index) উপাদানগুলিকে 0 থেকে 4 পর্যন্ত একটি scf.for লুপে নিষ্কাশন এবং সন্নিবেশ করতে ব্যবহৃত হয়। , tensor.extract এবং tensor.insert উভয়ই ভেক্টরাইজ করা যেতে পারে।
func.func @main(%input: tensor<12582912xbf16>, %output: tensor<12582912xbf16>) -> tensor<12582912xbf16> {
%vector_0 = arith.constant dense<0.000000e+00> : vector<4xbf16>
%0 = xla_gpu.apply_indexing #map(%thread_id_x, %block_id_x, %c0)
%2 = vector.transfer_read %input[%0], %cst {in_bounds = [true]} : tensor<12582912xbf16>, vector<4xbf16>
%xla_loop:2 = scf.for %vector_index = %c0 to %c4 step %c1
iter_args(%iter = %output, %iter_vector = %vector_0) -> (tensor<12582912xbf16>, vector<4xbf16>) {
%5 = vector.extract %2[%vector_index] : bf16 from vector<4xbf16>
%6 = arith.mulf %5, %5 : bf16
%7 = arith.mulf %6, %5 : bf16
%8 = arith.mulf %7, %cst_4 : bf16
%9 = arith.addf %5, %8 : bf16
%10 = arith.mulf %9, %cst_3 : bf16
%11 = math.tanh %10 : bf16
%12 = arith.addf %11, %cst_2 : bf16
%13 = arith.mulf %12, %cst_1 : bf16
%14 = arith.mulf %5, %13 : bf16
%15 = vector.insert %14, %iter_vector [%vector_index] : bf16 into vector<4xbf16>
scf.yield %iter, %15 : tensor<12582912xbf16>, vector<4xbf16>
}
%4 = vector.transfer_write %xla_loop#1, %output[%0] {in_bounds = [true]}
: vector<4xbf16>, tensor<12582912xbf16>
return %4 : tensor<12582912xbf16>
}
লুপ আনরোলিং
optimize_loops.cc দেখুন।
লুপ আনরোলিং scf.for লুপগুলি খুঁজে পায় যা আনরোল করা যেতে পারে। এই ক্ষেত্রে, ভেক্টরের উপাদানগুলির উপর লুপ অদৃশ্য হয়ে যায়।
func.func @main(%input: tensor<12582912xbf16>, %arg1: tensor<12582912xbf16>) -> tensor<12582912xbf16> {
%cst_0 = arith.constant dense<0.000000e+00> : vector<4xbf16>
%dim = xla_gpu.apply_indexing #map(%thread_id_x, %block_id_x, %c0)
%2 = vector.transfer_read %input[%dim], %cst {in_bounds = [true]} : tensor<12582912xbf16>, vector<4xbf16>
%3 = vector.extract %2[%c0] : bf16 from vector<4xbf16>
...
%13 = vector.insert %12, %cst_0 [%c0] : bf16 into vector<4xbf16>
%14 = vector.extract %2[%c1] : bf16 from vector<4xbf16>
...
%24 = vector.insert %23, %13 [%c1] : bf16 into vector<4xbf16>
%25 = vector.extract %2[%c2] : bf16 from vector<4xbf16>
...
%35 = vector.insert %34, %24 [%c2] : bf16 into vector<4xbf16>
%36 = vector.extract %2[%c3] : bf16 from vector<4xbf16>
...
%46 = vector.insert %45, %35 [%c3] : bf16 into vector<4xbf16>
%47 = vector.transfer_write %46, %arg1[%dim] {in_bounds = [true]} : vector<4xbf16>, tensor<12582912xbf16>
return %47 : tensor<12582912xbf16>
}
এলএলভিএম-এ রূপান্তর
আমরা বেশিরভাগই স্ট্যান্ডার্ড LLVM কম ব্যবহার করি, তবে কিছু বিশেষ পাস আছে। আমরা টেনসরের জন্য memref লোয়ারিং ব্যবহার করতে পারি না, যেহেতু আমরা IR বাফার করি না এবং আমাদের ABI memref ABI-এর সাথে সামঞ্জস্যপূর্ণ নয়। পরিবর্তে, আমাদের কাছে টেনসর থেকে সরাসরি LLVM এ একটি কাস্টম কমানো আছে।
- Lower_tensors.cc এ টেনসর কমানো হয়।
tensor.extractকমিয়েllvm.load,tensor.inserttollvm.store, স্পষ্টভাবে। - propagate_slice_indices এবং merge_pointers_to_same_slice একসাথে বাফার অ্যাসাইনমেন্ট এবং XLA এর ABI-এর বিস্তারিত প্রয়োগ করে: যদি দুটি টেনসর একই বাফার স্লাইস ভাগ করে, তবে সেগুলি একবার পাস করা হয়। এই পাসগুলি ফাংশন আর্গুমেন্টের অনুলিপি করে।
llvm.func @__nv_tanhf(f32) -> f32
llvm.func @main(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
%11 = nvvm.read.ptx.sreg.tid.x : i32
%12 = nvvm.read.ptx.sreg.ctaid.x : i32
%13 = llvm.mul %11, %1 : i32
%14 = llvm.mul %12, %0 : i32
%15 = llvm.add %13, %14 : i32
%16 = llvm.getelementptr inbounds %arg0[%15] : (!llvm.ptr, i32) -> !llvm.ptr, bf16
%17 = llvm.load %16 invariant : !llvm.ptr -> vector<4xbf16>
%18 = llvm.extractelement %17[%2 : i32] : vector<4xbf16>
%19 = llvm.fmul %18, %18 : bf16
%20 = llvm.fmul %19, %18 : bf16
%21 = llvm.fmul %20, %4 : bf16
%22 = llvm.fadd %18, %21 : bf16
%23 = llvm.fmul %22, %5 : bf16
%24 = llvm.fpext %23 : bf16 to f32
%25 = llvm.call @__nv_tanhf(%24) : (f32) -> f32
%26 = llvm.fptrunc %25 : f32 to bf16
%27 = llvm.fadd %26, %6 : bf16
%28 = llvm.fmul %27, %7 : bf16
%29 = llvm.fmul %18, %28 : bf16
%30 = llvm.insertelement %29, %8[%2 : i32] : vector<4xbf16>
...
}
ট্রান্সপোজ ইমিটার
আসুন একটু বেশি জড়িত উদাহরণ বিবেচনা করা যাক।
ট্রান্সপোজ ইমিটার লুপ ইমিটার থেকে আলাদা হয় শুধুমাত্র কিভাবে এন্ট্রি ফাংশন তৈরি হয়।
func.func @transpose(%arg0: tensor<20x160x170xf32>, %arg1: tensor<170x160x20xf32>) -> tensor<170x160x20xf32> {
%thread_id_x = gpu.thread_id x {xla.range = [0 : index, 127 : index]}
%block_id_x = gpu.block_id x {xla.range = [0 : index, 959 : index]}
%shmem = xla_gpu.allocate_shared : tensor<32x1x33xf32>
%xla_loop = xla_gpu.loop (%thread_id_x, %block_id_x)[%i, %j]
-> (%input_dim0, %input_dim1, %input_dim2, %shmem_dim0, %shmem_dim1, %shmem_dim2)
in #map iter_args(%iter = %shmem) -> (tensor<32x1x33xf32>) {
%extracted = tensor.extract %arg0[%input_dim0, %input_dim1, %input_dim2] : tensor<20x160x170xf32>
%0 = math.exp %extracted : f32
%inserted = tensor.insert %0 into %iter[%shmem_dim0, %shmem_dim1, %shmem_dim2] : tensor<32x1x33xf32>
xla_gpu.yield %inserted : tensor<32x1x33xf32>
}
%synced_tensor = xla_gpu.sync_threads %xla_loop : tensor<32x1x33xf32>
%xla_loop_0 = xla_gpu.loop (%thread_id_x %block_id_x)[%i, %j] -> (%dim0, %dim1, %dim2)
in #map1 iter_args(%iter = %arg1) -> (tensor<170x160x20xf32>) {
// indexing computations
%extracted = tensor.extract %synced_tensor[%0, %c0, %1] : tensor<32x1x33xf32>
%2 = math.absf %extracted : f32
%inserted = tensor.insert %2 into %iter[%3, %4, %1] : tensor<170x160x20xf32>
xla_gpu.yield %inserted : tensor<170x160x20xf32>
}
return %xla_loop_0 : tensor<170x160x20xf32>
}
এই ক্ষেত্রে, আমরা দুটি xla_gpu.loop ops তৈরি করি। প্রথমটি ইনপুট থেকে সমন্বিত পাঠ সম্পাদন করে এবং শেয়ার করা মেমরিতে ফলাফল লিখে।
শেয়ার্ড মেমরি টেনসর xla_gpu.allocate_shared op ব্যবহার করে তৈরি করা হয়েছে।
xla_gpu.sync_threads ব্যবহার করে থ্রেডগুলি সিঙ্ক্রোনাইজ করার পরে, দ্বিতীয় xla_gpu.loop শেয়ার করা মেমরি টেনসর থেকে উপাদানগুলি পড়ে এবং আউটপুটে একত্রিত লেখাগুলি সম্পাদন করে।
প্রজননকারী
কম্পাইলেশন পাইপলাইনের প্রতিটি পাসের পরে IR দেখার জন্য, কেউ --v=5 পতাকা সহ run_hlo_module চালু করতে পারে।
run_hlo_module --platform=CUDA --xla_disable_all_hlo_passes --reference_platform="" --v=5 /tmp/gelu.hlo
যেখানে /tmp/gelu.hlo রয়েছে
HloModule m:
gelu {
%param = bf16[6,512,4096] parameter(0)
%constant_0 = bf16[] constant(0.5)
%bcast_0 = bf16[6,512,4096] broadcast(bf16[] %constant_0), dimensions={}
%constant_1 = bf16[] constant(1)
%bcast_1 = bf16[6,512,4096] broadcast(bf16[] %constant_1), dimensions={}
%constant_2 = bf16[] constant(0.79785)
%bcast_2 = bf16[6,512,4096] broadcast(bf16[] %constant_2), dimensions={}
%constant_3 = bf16[] constant(0.044708)
%bcast_3 = bf16[6,512,4096] broadcast(bf16[] %constant_3), dimensions={}
%square = bf16[6,512,4096] multiply(bf16[6,512,4096] %param, bf16[6,512,4096] %param)
%cube = bf16[6,512,4096] multiply(bf16[6,512,4096] %square, bf16[6,512,4096] %param)
%multiply_3 = bf16[6,512,4096] multiply(bf16[6,512,4096] %cube, bf16[6,512,4096] %bcast_3)
%add_1 = bf16[6,512,4096] add(bf16[6,512,4096] %param, bf16[6,512,4096] %multiply_3)
%multiply_2 = bf16[6,512,4096] multiply(bf16[6,512,4096] %add_1, bf16[6,512,4096] %bcast_2)
%tanh_0 = bf16[6,512,4096] tanh(bf16[6,512,4096] %multiply_2)
%add_0 = bf16[6,512,4096] add(bf16[6,512,4096] %tanh_0, bf16[6,512,4096] %bcast_1)
%multiply_1 = bf16[6,512,4096] multiply(bf16[6,512,4096] %add_0, bf16[6,512,4096] %bcast_0)
ROOT %multiply_0 = bf16[6,512,4096] multiply(bf16[6,512,4096] %param, bf16[6,512,4096] %multiply_1)
}
ENTRY main {
%param = bf16[6,512,4096] parameter(0)
ROOT fusion = bf16[6,512,4096] fusion(%param), kind=kLoop, calls=gelu
}
কোড লিঙ্ক
- কম্পাইলেশন পাইপলাইন: mlir_fusion_emitter.h
- অপ্টিমাইজেশান এবং রূপান্তর পাস: gpu/fusions/transforms
- পার্টিশন লজিক: computation_partitioner.h
- হিরো-ভিত্তিক ইমিটার: জিপিইউ/ফিউশন
- XLA:GPU ops: xla_gpu_ops.td
- সঠিকতা এবং আলোক পরীক্ষা: জিপিইউ/ফিউশন/পরীক্ষা