این صفحه به‌وسیله ‏Cloud Translation API‏ ترجمه شده است.

XLA: امیترهای GPU

سه راه برای تولید کد برای HLO در XLA:GPU وجود دارد.

جایگزینی HLO با تماس های سفارشی به کتابخانه های خارجی، به عنوان مثال NVidia cuBLAS ، cuDNN .
کاشی کاری HLO به سطح بلوک و سپس استفاده از OpenAI Triton .
استفاده از XLA Emitters برای کاهش تدریجی HLO به LLVM IR.

این سند بر روی XLA:GPU Emitters متمرکز شده است.

کدژن مبتنی بر قهرمان

7 نوع امیتر در XLA:GPU وجود دارد. هر نوع امیتر مربوط به یک "قهرمان" فیوژن است، یعنی مهمترین عملیات در محاسبات ذوب شده که تولید کد را برای کل فیوژن شکل می دهد.

به عنوان مثال، اگر یک HloTransposeInstruction در ترکیب وجود داشته باشد که نیاز به استفاده از حافظه مشترک برای بهبود الگوهای خواندن و نوشتن حافظه دارد، فرستنده transpose انتخاب خواهد شد. ساطع کننده کاهش با استفاده از shuffles و حافظه مشترک کاهش ایجاد می کند. فرستنده حلقه امیتر پیش فرض است. اگر فیوژن هیرو نداشته باشد که برای آن امیتر خاصی داریم، از حلقه امیتر استفاده می شود.

نمای کلی سطح بالا

کد از بلوک های ساختمان بزرگ زیر تشکیل شده است:

پارتیشن‌کننده محاسبات - تقسیم محاسبات فیوژن HLO به توابع
Emitters - تبدیل HLO پارتیشن بندی شده به MLIR ( xla_gpu ، tensor ، arith ، math ، گویش scf )
خط لوله کامپایل - بهینه سازی و کاهش IR به LLVM

پارتیشن بندی

به computation_partitioner.h مراجعه کنید.

دستورالعمل های غیر عنصری HLO همیشه نمی توانند با هم منتشر شوند. نمودار HLO زیر را در نظر بگیرید:

     param
       |
      log
      |  \
      |  transpose
      |  /
      add

اگر این را در یک تابع واحد منتشر کنیم، log در دو شاخص مختلف برای هر عنصر add قابل دسترسی خواهد بود. قطره چکان های قدیمی این مشکل را با دو بار تولید log حل می کنند. برای این نمودار خاص، این یک مشکل نیست، اما زمانی که چندین تقسیم وجود دارد، اندازه کد به طور تصاعدی رشد می کند.

در اینجا، ما این مشکل را با پارتیشن بندی نمودار به قطعاتی که می توانند به صورت ایمن به عنوان یک تابع منتشر شوند، حل می کنیم. معیارها عبارتند از:

دستورالعمل هایی که فقط یک کاربر دارند می توانند همراه با کاربر خود منتشر شوند.
دستورالعمل‌هایی که چندین کاربر دارند، در صورتی که همه کاربران از طریق شاخص‌های یکسانی به آن‌ها دسترسی داشته باشند، همراه با کاربران خود ارسال می‌شوند.

در مثال بالا، add و tranpose به شاخص‌های مختلفی از log دسترسی دارند، بنابراین انتشار آن همراه با آنها امن نیست.

بنابراین گراف به سه تابع تقسیم می‌شود (که هر کدام فقط یک دستورالعمل دارد).

همین امر در مورد مثال زیر با slice و pad add نیز صدق می کند.

انتشار عنصری

elemental_hlo_to_mlir.h را ببینید.

انتشار عنصری حلقه ها و عملیات ریاضی/حساب را برای HloInstructions ایجاد می کند. در بیشتر موارد، این ساده است، اما چیزهای جالبی در اینجا وجود دارد.

تحولات نمایه سازی

برخی از دستورالعمل‌ها ( transpose ، broadcast ، reshape ، slice ، reverse و چند مورد دیگر) صرفاً تغییراتی در شاخص‌ها هستند: برای تولید یک عنصر از نتیجه، ما نیاز به تولید عنصر دیگری از ورودی داریم. برای این کار، می‌توانیم از indexing_analysis XLA استفاده کنیم، که دارای توابعی برای تولید خروجی به نگاشت ورودی برای یک دستورالعمل است.

به عنوان مثال، برای transpose از [20,40] به [40,20] ، نقشه نمایه سازی زیر را ایجاد می کند (یک عبارت افین برای هر بعد ورودی؛ d0 و d1 ابعاد خروجی هستند):

  (d0, d1) -> d1
  (d0, d1) -> d0

بنابراین برای این دستورالعمل‌های تبدیل شاخص خالص، می‌توانیم به سادگی نقشه را دریافت کنیم، آن را روی شاخص‌های خروجی اعمال کنیم و ورودی را در شاخص حاصل تولید کنیم.

به طور مشابه، pad op از نقشه‌های نمایه‌سازی و محدودیت‌ها برای بیشتر پیاده‌سازی استفاده می‌کند. pad همچنین یک تبدیل نمایه سازی با چند بررسی اضافه شده است تا ببینیم آیا عنصری از ورودی یا مقدار padding را برمی گردانیم.

تاپل ها

ما از 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 را به‌روزرسانی می‌کند و سپس آنها را برمی‌گرداند. هیچ استفاده دیگری از تانسورهای خروجی مجاز نیست.

خط لوله تدوین

ساطع کننده حلقه

به loop_mlir.h مراجعه کنید.

بیایید مهم ترین پاس های خط لوله کامپایل MLIR را با استفاده از HLO برای تابع GELU مطالعه کنیم.

این محاسبات HLO فقط دارای عملیات عنصری، ثابت و پخش است. با استفاده از امیتر حلقه منتشر می شود.

تبدیل 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`

low_xla_gpu_to_scf.cc را ببینید.

xla_gpu.loop نشان دهنده یک لانه حلقه با یک بررسی مرزی در داخل است. اگر متغیرهای القایی حلقه خارج از محدوده دامنه نقشه نمایه سازی باشند، این تکرار نادیده گرفته می شود. به این معنی که حلقه به 1 یا چند عملیات تو در تو scf.for با یک 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 روی 1 بعدی پیش بینی می شوند. این امر بردارسازی و کاهش به 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 تجزیه و تحلیل می‌کند و اگر آنها توسط 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) برای محاسبه عناصر برای استخراج و درج در یک حلقه scf.for از 0 تا 4 استفاده می شود. بنابراین ، هم 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 مراجعه کنید.

حلقه unrolling حلقه‌های 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

ما بیشتر از پایین آوردن استاندارد LLVM استفاده می کنیم، اما چند پاس ویژه وجود دارد. ما نمی توانیم از کاهش های memref برای تانسورها استفاده کنیم، زیرا ما IR را بافر نمی کنیم و ABI ما با memref ABI سازگار نیست. در عوض، ما یک کاهش سفارشی مستقیماً از تانسورها به LLVM داریم.

کاهش تانسورها در low_tensors.cc انجام می شود. tensor.extract به روش واضح به llvm.load ، tensor.insert به llvm.store کاهش می یابد.
propagate_slice_indices و merge_pointers_to_same_slice با هم جزئیاتی از انتساب بافر و ABI XLA را پیاده سازی می کنند: اگر دو تانسور یک برش بافر را به اشتراک بگذارند، فقط یک بار ارسال می شوند. این پاس ها آرگومان های تابع را کپی می کنند.

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

انتقال امیتر

بیایید یک مثال کمی درگیرتر را در نظر بگیریم.

امیتر transpose با امیتر حلقه فقط در نحوه تولید تابع ورودی متفاوت است.

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 تولید می کنیم. اولین مورد خواندن های ترکیبی را از ورودی انجام می دهد و نتیجه را در حافظه مشترک می نویسد.

تانسور حافظه مشترک با استفاده از xla_gpu.allocate_shared op ایجاد می شود.

پس از اینکه رشته ها با استفاده از xla_gpu.sync_threads همگام شدند، xla_gpu.loop دوم عناصر را از تانسور حافظه مشترک می خواند و نوشتن های ترکیبی را در خروجی انجام می دهد.

بازتولید کننده

برای دیدن IR پس از هر عبور از خط لوله کامپایل، می توان run_hlo_module با پرچم --v=5 راه اندازی کرد.

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
امیترهای مبتنی بر قهرمان: gpu/fusions
XLA: عملیات گرافیکی: xla_gpu_ops.td
تست های صحت و روشن: gpu/fusions/tests