XLA Özel Çağrılar
Bu belgede, XLA "özel çağrılarının" nasıl yazılacağı ve kullanılacağı açıklanmaktadır. Özel çağrılar, bir XLA programından C ++ veya CUDA gibi bir programlama dilinde yazılmış kodu çağırmanızı sağlar.
Uyarı: Özel aramalar düşük seviyeli uzman kullanıcı özelliğidir. Özel aramaları kullanarak programınızı hata ayıklaması zor (ve hatta fark edilmesi zor) yollarla kırmak kolaydır. Bir şeyler ters gittiğinde XLA'da hata ayıklamaya hazır olmadığınız sürece özel çağrıları kullanmamalısınız ve sorun yaşarsanız XLA geliştiricilerinden nispeten daha az yardım beklemelisiniz.
Uyarı: Özel çağrı API / ABI şu anda kararlı değil. Bunu kaprisli bir şekilde değiştirmek istemiyoruz, ama değişebilir. Gelecekteki olası bazı değişiklikler aşağıda açıklanmaktadır.
CPU'da özel çağrı
XLA'nın istemci API'sı aracılığıyla özel bir çağrıyı temsil eden bir HLO talimatı oluşturabilirsiniz. Bu, yazılı olarak TensorFlow aracılığıyla gösterilmez.
Örneğin, aşağıdaki kod CPU'da A[i] = B[i % 128] + C[i] hesaplamak için özel bir çağrı kullanır. (Elbette bunu düzenli HLO ile yapabilirsin ve yapmalısın!)
do_custom_call fonksiyonunun üzerinde çalıştığı tamponların boyutlarını do_custom_call gerektiğine dikkat edin. Bu örnekte 128 ve 2048 boyutlarını sabit kodluyoruz. Bunu yapmak istemiyorsanız, boyutları aramaya parametre olarak iletebilirsiniz.
GPU'da özel çağrı
GPU özel çağrı çerçevesi CPU'dekinden biraz farklıdır. Yukarıdaki CPU kodu ile aynı A[i] = B[i % 128] + C[i] hesaplamasını yapan bir CUDA örneği.
İlk olarak GPU özel çağrı işlevinin hala CPU'da yürütülen bir işlev olduğuna dikkat edin. do_custom_call CPU fonksiyonumuz GPU üzerinde çalışmayı sıralamaktan sorumludur. Burada bir CUDA çekirdeği piyasaya sürüyor, ancak çağrı küpleri gibi başka bir şey de yapabilir.
buffers , ana bilgisayarda yaşayan bir işaretçiler dizisidir ve içerdiği her öğe, aygıt (yani GPU) belleğine puan verir. Parametreler önce gelir, ardından çıktı değeri gelir. Bu iki parametreler, vardır CPU çağrı kuralı, oldukça farklıdır ins ve out . Ayırmamızın ana nedeni, tuple şeklindeki girişlerin / çıkışların verimli bir şekilde kullanılmasını mümkün kılmaktır; aşağıdaki bölüme bakın.
CPU örneğinde olduğu gibi, giriş ve çıkış arabellek boyutlarını özel aramamıza sabit olarak kodladık. Ancak CPU durumundan farklı olarak, arabellek boyutlarını işlenen olarak özel aramaya geçirmek iyi çalışmaz. Genellikle CPU'da kullanabileceğimiz tampon boyutlarına ihtiyacımız var; örneğin bir çekirdeği başlatırken, kullanılacak blok / ızgara boyutlarını bilmemiz gerekir. Ancak arabellek boyutlarını işlenen olarak özel çağrımıza geçirecek olsaydık, değerleri GPU belleğinde yaşayacaktı. Daha sonra, sadece boyutları okumak için operasyonumuzun başında pahalı bir senkronize cihazdan ana bilgisayara memcpy yapmamız gerekir.
Bu sorunu çözmenize izin vermek için opaque parametreyi sağlıyoruz. Özel aramayı oluştururken bunu keyfi bir bayt dizesine ayarlayabilirsiniz:
xla::Shape bir protokol tamponu temsiline sahip olduğundan, bu serileştirilmiş protokolü opaque içinde saklayabilir ve GPU özel çağrınızda serisini kaldırabilirsiniz. Ancak xla::ShapeProto sık sık değişmese de değiştiğine dikkat edin. Geçmişte nasıl değiştiğini görmek için git günlüğünü kontrol edin.
Özel çağrılara tuples aktarma
Aşağıdaki özel aramayı düşünün.
Hem CPU hem de GPU'da, bir demet bellekte bir işaretçi dizisi olarak temsil edilir. C ++ - sözde kodunda, yukarıdaki parametre 0 aşağıdaki gibi düzenlenir.
Her ne kadar tupleslerin bellek içi gösterimi CPU ve GPU'da aynı olsa da, CPU ve GPU özel çağrı arama kurallarında farklı şekilde ele alınır.
Geçici tamponlar olarak grup çıkışları
Özel çağrılara grup girişleri kolaylık sağlar, ancak kesinlikle gerekli değildir. Özel aramalarda grup girişlerini desteklemiyorsak, özel aramaya geçmeden önce tuples'ı get-tuple-element kullanarak açabilirsiniz.
Öte yandan, grup çıktıları başka türlü yapamayacağınız şeyleri yapmanızı sağlar.
Grup çıkışlarına sahip olmanın bariz nedeni, özel bir aramanın (veya başka bir XLA opunun) çoklu bağımsız dizileri bu şekilde döndürmesidir.
Ancak daha az belli ki, bir grup çıkışı aynı zamanda özel çağrı temp belleğinizi vermenin bir yoludur. Evet, bir çıktı geçici arabelleği temsil edebilir. Bir çıkış arabelleğinin, op'un ona yazabileceği özelliğe sahip olduğunu ve yazıldıktan sonra ondan okuyabileceğini düşünün. Geçici bir arabellekten tam olarak istediğiniz budur.
Yukarıdaki örnekte, F32 i geçici bir tampon olarak kullanmak istediğimizi varsayalım. Sonra HLO'yu yukarıdaki gibi yazardık ve özel çağrının çıktısının 1 numaralı dizinini asla okumazdık.
CPU özel çağrılarında tuples
CPU kodunda, do_custom_call(const void** ins, void* out) . ins , param0 işaret eden tek bir öğeye sahip bir dizidir. Arasında subbuffers param0 bu işaretçiyi kaldırma tarafından erişilebilir ve subbuffers output_tuple kaldırma tarafından erişilebilir out .
GPU özel çağrılarındaki tuples
GPU kodunda, do_custom_call(..., void** buffers, ...) . Bu durumda buffers , giriş / çıkıştaki her bir yaprak arabelleği için bir tane olmak üzere altı aygıt işaretçisi içeren bir ana bilgisayar dizisidir. Düz listeyi oluşturmak için, parametreler ve çıktılar üzerinde tekrarlar ve her biri için şeklinin ön geçişini yaparız. somut olarak: