XLA和PyTorch的鏈接

前言

XLA (Accelerated Linear Algebra)是一個(gè)開源的機(jī)器學(xué)習(xí)編譯器，對(duì)PyTorch、Tensorflow、JAX等多個(gè)深度學(xué)習(xí)框架都有支持。最初XLA實(shí)際上是跟Tensorflow深度結(jié)合的，很好地服務(wù)了Tensorflow和TPU，而與XLA的結(jié)合主要依賴于社區(qū)的支持，即torch-xla。

torch-xla在支持XLA編譯的基礎(chǔ)上，較大限度地保持了PyTorch的易用性，貼一個(gè)官方的DDP訓(xùn)練的例子：

importtorch.distributedasdist
-importtorch.multiprocessingasmp
+importtorch_xla.core.xla_modelasxm
+importtorch_xla.distributed.parallel_loaderaspl
+importtorch_xla.distributed.xla_multiprocessingasxmp
+importtorch_xla.distributed.xla_backend

def_mp_fn(rank,world_size):
...

-os.environ['MASTER_ADDR']='localhost'
-os.environ['MASTER_PORT']='12355'
-dist.init_process_group("gloo",rank=rank,world_size=world_size)
+#RankandworldsizeareinferredfromtheXLAdeviceruntime
+dist.init_process_group("xla",init_method='xla://')
+
+model.to(xm.xla_device())
+#`gradient_as_bucket_view=True`requiredforXLA
+ddp_model=DDP(model,gradient_as_bucket_view=True)

-model=model.to(rank)
-ddp_model=DDP(model,device_ids=[rank])
+xla_train_loader=pl.MpDeviceLoader(train_loader,xm.xla_device())

-forinputs,labelsintrain_loader:
+forinputs,labelsinxla_train_loader:
optimizer.zero_grad()
outputs=ddp_model(inputs)
loss=loss_fn(outputs,labels)
loss.backward()
optimizer.step()

if__name__=='__main__':
-mp.spawn(_mp_fn,args=(),nprocs=world_size)
+xmp.spawn(_mp_fn,args=())

將一段PyTorch代碼改寫為torch-xla代碼，主要就是三個(gè)方面：

將模型和數(shù)據(jù)放到xla device上

適當(dāng)?shù)臅r(shí)候調(diào)用xm.mark_step

某些組件該用pytorchx-xla提供的，比如amp和spawn

其中第二條并沒有在上面的代碼中體現(xiàn)，原因是為了讓用戶少改代碼，torch-xla將mark_step封裝到了dataloader中，實(shí)際上不考慮DDP的完整訓(xùn)練的過程可以簡(jiǎn)寫如下：

device=xm.xla_device()
model=model.to(device)
fordata,labelinenumerate(dataloader):
data,label=data.to(device),label.to(device)
output=model(data)
loss=func(output,label)
loss.backward()
optimizer.step()
xm.mark_step()

xm.mark_step的作用就是"告訴"框架：現(xiàn)在對(duì)圖的定義告一段落了，可以編譯并執(zhí)行計(jì)算了。既然如此，那么mark_step之前的內(nèi)容是做了什么呢？因?yàn)橐趍ark_step之后才編譯并計(jì)算，那么前面肯定不能執(zhí)行實(shí)際的運(yùn)算。這就引出了Trace和LazyTensor的概念。

其實(shí)到了這里，如果對(duì)tensorflow或者torch.fx等比較熟悉，就已經(jīng)很容易理解了，在mark_step之前，torch-xla將torch Tensor換成了LazyTensor，進(jìn)而將原本是PyTorch中eager computation的過程替換成了trace的過程，最后生成一張計(jì)算圖來優(yōu)化和執(zhí)行。簡(jiǎn)而言之這個(gè)過程是PyTorch Tensor -> XLATensor -> HLO IR，其中HLO就是XLA所使用的IR。在每次調(diào)用到torch op的時(shí)候，會(huì)調(diào)用一次GetIrValue，這時(shí)候就意味著一個(gè)節(jié)點(diǎn)被寫入了圖中。更具體的信息可以參考XLA Tensor Deep Dive這部分文檔。需要注意的是，trace這個(gè)過程是獨(dú)立于mark_step的，即便你的每個(gè)循環(huán)都不寫mark_step，這個(gè)循環(huán)也可以一直持續(xù)下去，只不過在這種情況下，永遠(yuǎn)都不會(huì)發(fā)生圖的編譯和執(zhí)行，除非在某一步trace的時(shí)候，發(fā)現(xiàn)圖的大小已經(jīng)超出了pytorch-xla允許的上限。

PyTorch與torch-xla的橋接

知曉了Trace過程之后，就會(huì)好奇一個(gè)問題：當(dāng)用戶執(zhí)行一個(gè)PyTorch函數(shù)調(diào)用的時(shí)候，torch-xla怎么將這個(gè)函數(shù)記錄下來的？

最容易想到的答案是“torch-xla作為PyTorch的一個(gè)編譯選項(xiàng)，打開的時(shí)候就會(huì)使得二者建立起映射關(guān)系”，但很可惜，這個(gè)答案是錯(cuò)誤的，仔細(xì)看PyTorch的CMake文件以及torch-xla的編譯方式就會(huì)明白，torch-xla是幾乎單向依賴于PyTorch的（為什么不是全部后面會(huì)講）。既然PyTorch本身在編譯期間并不知道torch-xla的存在，那么當(dāng)用戶使用一個(gè)xla device上的Tensor作為一個(gè)torch function的輸入的時(shí)候，又經(jīng)歷了怎樣一個(gè)過程調(diào)用到pytorch-xla中的東西呢？

從XLATensor開始的溯源

盡管我們現(xiàn)在并不知道怎么調(diào)用到torch-xla中的，但我們知道PyTorch Tensor一定要轉(zhuǎn)換成XLATensor（參考tensor.h），那么我們只需要在關(guān)鍵的轉(zhuǎn)換之處打印出調(diào)用堆棧，自然就可以找到調(diào)用方，這樣雖然不能保證找到PyTorch中的位置，但是能夠找到torch-xla中最上層的調(diào)用。注意到XLATensor只有下面這一個(gè)創(chuàng)建函數(shù)接受at::Tensor作為輸入，因此就在這里面打印調(diào)用棧。

XLATensorXLATensor::Create(constat::Tensor&tensor,constDevice&device)

測(cè)試的用例很簡(jiǎn)單，我們讓兩個(gè)xla device上的Tensor相乘：

importtorch_xla.core.xla_modelasxm
importtorch

device=xm.xla_device()
a=torch.normal(0,1,(2,3)).to(device)
b=torch.normal(0,1,(2,3)).to(device)

c=a*b

在上述位置插入堆棧打印代碼并重新編譯、安裝后運(yùn)行用例，可以看到以下輸出（截取部分）：

usr/local/lib/python3.8/dist-packages/_XLAC.cpython-38-x86_64-linux-gnu.so(_ZN9torch_xla15TensorToXlaDataERKN2at6TensorERKNS_6DeviceEb+0x64d)[0x7f086098b9ed]
/usr/local/lib/python3.8/dist-packages/_XLAC.cpython-38-x86_64-linux-gnu.so(_ZNK9torch_xla9XLATensor19GetIrValueForTensorERKN2at6TensorERKNS_6DeviceE+0xa5)[0x7f0860853955]
/usr/local/lib/python3.8/dist-packages/_XLAC.cpython-38-x86_64-linux-gnu.so(_ZNK9torch_xla9XLATensor10GetIrValueEv+0x19b)[0x7f0860853d5b]
/usr/local/lib/python3.8/dist-packages/_XLAC.cpython-38-x86_64-linux-gnu.so(_ZN9torch_xla9XLATensor3mulERKS0_S2_N3c108optionalINS3_10ScalarTypeEEE+0x3f)[0x7f086087631f]
/usr/local/lib/python3.8/dist-packages/_XLAC.cpython-38-x86_64-linux-gnu.so(_ZN9torch_xla18XLANativeFunctions3mulERKN2at6TensorES4_+0xc4)[0x7f08606d4da4]
/usr/local/lib/python3.8/dist-packages/_XLAC.cpython-38-x86_64-linux-gnu.so(+0x19d158)[0x7f08605f7158]
/usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so(_ZN2at4_ops10mul_Tensor10redispatchEN3c1014DispatchKeySetERKNS_6TensorES6_+0xc5)[0x7f0945c9d055]
/usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so(+0x2b8986c)[0x7f094705986c]
/usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so(+0x2b8a37b)[0x7f094705a37b]
/usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so(_ZN2at4_ops10mul_Tensor4callERKNS_6TensorES4_+0x157)[0x7f0945cee717]
/usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_python.so(+0x3ee91f)[0x7f094e4b391f]
/usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_python.so(+0x3eeafb)[0x7f094e4b3afb]
python()[0x5042f9]

明顯可以看到是從python的堆棧調(diào)用過來的，分析一下可以得知_ZN2at4_ops10mul_Tensor10redispatchEN3c1014DispatchKeySetERKNS_6TensorES6_+0xc5對(duì)應(yīng)的定義是at::DispatchKeySet, at::Tensor const&, at::Tensor const&)+0xc5

雖然這里意義仍有些不明，但我們已經(jīng)可以做出推測(cè)了：redistpatch函數(shù)是根據(jù)DispatchKeySet來決定將操作dispatch到某個(gè)backend上，xla的device信息就被包含在其中。而后面兩個(gè)輸入的const at::Tensor&就是乘法操作的兩個(gè)輸入。

根據(jù)上面的關(guān)鍵字redispatch來尋找，我們可以找到這樣一個(gè)文件gen.py，其中的codegen函數(shù)很多，但最顯眼的是下面的OperatorGen：

@dataclass(frozen=True)
classComputeOperators:
target:Union[
Literal[Target.DECLARATION],
Literal[Target.DEFINITION]
]

@method_with_native_function
def__call__(self,f:NativeFunction)->str:
sig=DispatcherSignature.from_schema(f.func)
name=f.func.name.unambiguous_name()
call_method_name='call'
redispatch_method_name='redispatch'

ifself.targetisTarget.DECLARATION:
returnf"""
structTORCH_API{name}{{
usingschema={sig.type()};
usingptr_schema=schema*;
//SeeNote[staticconstexprchar*membersforwindowsNVCC]
STATIC_CONSTEXPR_STR_INL_EXCEPT_WIN_CUDA(name,"aten::{f.func.name.name}")
STATIC_CONSTEXPR_STR_INL_EXCEPT_WIN_CUDA(overload_name,"{f.func.name.overload_name}")
STATIC_CONSTEXPR_STR_INL_EXCEPT_WIN_CUDA(schema_str,{cpp_string(str(f.func))})
static{sig.defn(name=call_method_name,is_redispatching_fn=False)};
static{sig.defn(name=redispatch_method_name,is_redispatching_fn=True)};
}};"""
elifself.targetisTarget.DEFINITION:
defns=f"""
STATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name},name,"aten::{f.func.name.name}")
STATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name},overload_name,"{f.func.name.overload_name}")
STATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name},schema_str,{cpp_string(str(f.func))})

//aten::{f.func}
staticC10_NOINLINEc10::TypedOperatorHandle<{name}::schema>create_{name}_typed_handle(){{
returnc10::singleton()
.findSchemaOrThrow({name}::name,{name}::overload_name)
.typed<{name}::schema>();
}}
"""

foris_redispatching_fnin[False,True]:
ifis_redispatching_fn:
dispatcher_exprs_str=','.join(['dispatchKeySet']+[a.nameforainsig.arguments()])
dispatcher_call='redispatch'
method_name=f'{name}::{redispatch_method_name}'
else:
dispatcher_exprs_str=','.join([a.nameforainsig.arguments()])
dispatcher_call='call'
method_name=f'{name}::{call_method_name}'

defns+=f"""
//aten::{f.func}
{sig.defn(name=method_name,is_redispatching_fn=is_redispatching_fn)}{{
staticautoop=create_{name}_typed_handle();
returnop.{dispatcher_call}({dispatcher_exprs_str});
}}
"""
returndefns
else:
assert_never(self.target)

對(duì)于每個(gè)算子，PyTorch會(huì)（在編譯前）在這里生成許多類，這些類會(huì)有靜態(tài)成員call或者redispatch，其中redispatch負(fù)責(zé)分發(fā)具體的實(shí)現(xiàn)。這里的codegen比較繁瑣，這里就不再細(xì)講。

注冊(cè)PyTorch庫實(shí)現(xiàn)

即便我們找到了上面redispatch和codegen的線索，看起來仍然不足以解釋PyTorch到torch-xla的橋接，因?yàn)镻yTorch和torch-xla兩個(gè)庫之間的調(diào)用，必須要有符號(hào)的映射才可以，而不是一些函數(shù)形式上的相同。PyTorch中是有Dispatcher機(jī)制的，這個(gè)機(jī)制很常見于很多框架，比如oneflow也是有一套類似的Dispatcher機(jī)制。這套機(jī)制最大的好處就是在盡可能減少侵入式修改的前提下保證了較高的可擴(kuò)展性。簡(jiǎn)而言之，我們的op有一種定義，但可以有多種實(shí)現(xiàn)方式，并且這個(gè)實(shí)現(xiàn)的代碼可以不在框架內(nèi)部，這樣就使得框架在保持通用性的同時(shí)，易于在特定環(huán)境下做針對(duì)性的擴(kuò)展。這套機(jī)制本質(zhì)上就是建立了一個(gè)字典，將op映射到函數(shù)指針，那么每次調(diào)用一個(gè)op的時(shí)候，我們可以根據(jù)一些標(biāo)識(shí)（比如tensor.device）來判斷應(yīng)該調(diào)用哪一種實(shí)現(xiàn)。

PyTorch中提供了一個(gè)宏用來將實(shí)現(xiàn)注冊(cè)，從而讓dispatcher可以調(diào)用：

#define_TORCH_LIBRARY_IMPL(ns,k,m,uid)
staticvoidC10_CONCATENATE(
TORCH_LIBRARY_IMPL_init_##ns##_##k##_,uid)(torch::Library&);
staticconsttorch::TorchLibraryInitC10_CONCATENATE(
TORCH_LIBRARY_IMPL_static_init_##ns##_##k##_,uid)(
torch::IMPL,
c10::if_constexpr(
[](){
return&C10_CONCATENATE(
TORCH_LIBRARY_IMPL_init_##ns##_##k##_,uid);
},
[](){return[](torch::Library&)->void{};}),
#ns,
c10::k),
__FILE__,
__LINE__);
voidC10_CONCATENATE(
TORCH_LIBRARY_IMPL_init_##ns##_##k##_,uid)(torch::Library&m)

這個(gè)宏如果完全展開會(huì)是下面這樣：

staticvoidTORCH_LIBRARY_IMPL_init_aten_CPU_0(torch::Library&);
staticconsttorch::detail::TorchLibraryInitTORCH_LIBRARY_IMPL_static_init_aten_CPU_0(
torch::IMPL,
(c10::CPU)
?&TORCH_LIBRARY_IMPL_init_aten_CPU_0
:[](torch::Library&)->void{}),
"aten",
c10::CPU),
__FILE__,
__LINE__);
voidTORCH_LIBRARY_IMPL_init_aten_CPU_0(torch::Library&m)

這里比較需要注意的是第二行的TORCH_LIBRARY_IMPL_static_init_aten_CPU_0并不是一個(gè)函數(shù)，而是一個(gè)靜態(tài)變量，它的作用就是在torch_xla庫初始化的時(shí)候，將xla定義的op注冊(cè)到PyTorch中。

關(guān)于這部分更詳細(xì)的介紹可以參考https://zhuanlan.zhihu.com/p/648578629。

從PyTorch調(diào)用到torch_xla

xla調(diào)用上面所說的宏進(jìn)行注冊(cè)的位置在RegisterXLA.cpp這個(gè)文件中（codegen的結(jié)果），如下：

ORCH_LIBRARY_IMPL(aten,XLA,m){
m.impl("abs",
TORCH_FN(wrapper__abs));

...
}

其中，wrapper__abs的定義如下：

at::Tensorwrapper__abs(constat::Tensor&self){
returntorch_xla::abs(self);
}

顯然，這個(gè)定義和PyTorch框架內(nèi)部的算子是完全一致的，只是修改了實(shí)現(xiàn)。而XLANativeFunctions::abs的實(shí)現(xiàn)可以在aten_xla_type.cpp中找到，如下所示：

at::TensorXLANativeFunctions::abs(constat::Tensor&self){
XLA_FN_COUNTER("xla::");
returnbridge::abs(bridge::GetXlaTensor(self)));
}

到這里已經(jīng)比較明朗了，注冊(cè)之后，PyTorch上對(duì)于op的調(diào)用最終會(huì)進(jìn)入torch_xla的native function中調(diào)用對(duì)應(yīng)的op實(shí)現(xiàn)，而這些實(shí)現(xiàn)的根本都是對(duì)XLATensor進(jìn)行操作，在最終操作執(zhí)行完成之后，會(huì)將作為結(jié)果的XLATensor重新轉(zhuǎn)換為torch Tensor，但要注意，這里的結(jié)果不一定被實(shí)際計(jì)算了，也可能只是記錄了一下IR，將節(jié)點(diǎn)加入圖中，這取決于具體的實(shí)現(xiàn)。

總結(jié)

其實(shí)torch-xla官方的文檔里是有關(guān)于代碼生成和算子注冊(cè)這個(gè)過程的描述的，只不過一開始我沒找到這個(gè)文檔，走了一點(diǎn)彎路，但是自己探索也會(huì)覺得更明了這個(gè)過程。官方文檔中的描述如下（節(jié)選）：

All file mentioned below lives under the xla/torch_xla/csrc folder, with the exception of codegen/xla_native_functions.yaml

xla_native_functions.yaml contains the list of all operators that are lowered. Each operator name must directly match a pytorch operator listed in native_functions.yaml. This file serves as the interface to adding new xla operators, and is an input to PyTorch's codegen machinery. It generates the below 3 files: XLANativeFunctions.h, RegisterXLA.cpp, and RegisterAutogradXLA.cpp

XLANativeFunctions.h and aten_xla_type.cpp are entry points of PyTorch to the pytorch_xla world, and contain the manually written lowerings to XLA for each operator. XLANativeFunctions.h is auto-generated through a combination of xla_native_functions.yaml and the PyTorch core native_functions.yaml file, and contains declarations for kernels that need to be defined in aten_xla_type.cpp. The kernels written here need to construct 'XLATensor' using the input at::Tensor and other parameters. The resulting XLATensor needs to be converted back to the at::Tensor before returning to the PyTorch world.

RegisterXLA.cpp and RegisterAutogradXLA.cpp are auto-generated files that register all lowerings to the PyTorch Dispatcher. They also include auto-generated wrapper implementations of out= and inplace operators.

大概意思就是實(shí)際上torch-xla就是根據(jù)xla_native_functions.yaml這個(gè)文件來生成算子的定義，然后再生成對(duì)應(yīng)的RegisterXLA.cpp中的注冊(cè)代碼，這也跟PyTorch的codegen方式一致。

綜合這一整個(gè)過程可以看出，PyTorch是保持了高度的可擴(kuò)展性的，不需要多少侵入式的修改就可以將所有的算子全部替換成自己的，這樣的方式也可以讓開發(fā)者不用去關(guān)注dispatcher及其上層的實(shí)現(xiàn)，專注于算子本身的邏輯。

編輯：黃飛