XLA和PyTorch的鏈接

前言

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

torch-xla在支持XLA編譯的基礎(chǔ)上，較大限度地保持了PyTorch的易用性，貼一個官方的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代碼，主要就是三個方面：

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

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

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

其中第二條并沒有在上面的代碼中體現(xiàn)，原因是為了讓用戶少改代碼，torch-xla將mark_step封裝到了dataloader中，實際上不考慮DDP的完整訓(xù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)在對圖的定義告一段落了，可以編譯并執(zhí)行計算了。既然如此，那么mark_step之前的內(nèi)容是做了什么呢？因為要在mark_step之后才編譯并計算，那么前面肯定不能執(zhí)行實際的運算。這就引出了Trace和LazyTensor的概念。

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

PyTorch與torch-xla的橋接

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

最容易想到的答案是“torch-xla作為PyTorch的一個編譯選項，打開的時候就會使得二者建立起映射關(guān)系”，但很可惜，這個答案是錯誤的，仔細看PyTorch的CMake文件以及torch-xla的編譯方式就會明白，torch-xla是幾乎單向依賴于PyTorch的（為什么不是全部后面會講）。既然PyTorch本身在編譯期間并不知道torch-xla的存在，那么當用戶使用一個xla device上的Tensor作為一個torch function的輸入的時候，又經(jīng)歷了怎樣一個過程調(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只有下面這一個創(chuàng)建函數(shù)接受at::Tensor作為輸入，因此就在這里面打印調(diào)用棧。

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

測試的用例很簡單，我們讓兩個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

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

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對應(yīng)的定義是at::DispatchKeySet, at::Tensor const&, at::Tensor const&)+0xc5

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

根據(jù)上面的關(guān)鍵字redispatch來尋找，我們可以找到這樣一個文件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)

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

注冊PyTorch庫實現(xiàn)

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

PyTorch中提供了一個宏用來將實現(xiàn)注冊，從而讓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)

這個宏如果完全展開會是下面這樣：

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并不是一個函數(shù)，而是一個靜態(tài)變量，它的作用就是在torch_xla庫初始化的時候，將xla定義的op注冊到PyTorch中。

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

從PyTorch調(diào)用到torch_xla

xla調(diào)用上面所說的宏進行注冊的位置在RegisterXLA.cpp這個文件中（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);
}

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

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

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

總結(jié)

其實torch-xla官方的文檔里是有關(guān)于代碼生成和算子注冊這個過程的描述的，只不過一開始我沒找到這個文檔，走了一點彎路，但是自己探索也會覺得更明了這個過程。官方文檔中的描述如下（節(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.

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

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

編輯：黃飛