XLA和PyTorch的鏈接

前言

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

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

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代碼，主要就是三個方面：

將模型和數據放到xla device上

適當的時候調用xm.mark_step

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

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

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的作用就是"告訴"框架：現在對圖的定義告一段落了，可以編譯并執行計算了。既然如此，那么mark_step之前的內容是做了什么呢？因為要在mark_step之后才編譯并計算，那么前面肯定不能執行實際的運算。這就引出了Trace和LazyTensor的概念。

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

PyTorch與torch-xla的橋接

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

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

從XLATensor開始的溯源

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

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的堆棧調用過來的，分析一下可以得知_ZN2at4_ops10mul_Tensor10redispatchEN3c1014DispatchKeySetERKNS_6TensorES6_+0xc5對應的定義是at::DispatchKeySet, at::Tensor const&, at::Tensor const&)+0xc5

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

根據上面的關鍵字redispatch來尋找，我們可以找到這樣一個文件gen.py，其中的codegen函數很多，但最顯眼的是下面的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會（在編譯前）在這里生成許多類，這些類會有靜態成員call或者redispatch，其中redispatch負責分發具體的實現。這里的codegen比較繁瑣，這里就不再細講。

注冊PyTorch庫實現

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

PyTorch中提供了一個宏用來將實現注冊，從而讓dispatcher可以調用：

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

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

從PyTorch調用到torch_xla

xla調用上面所說的宏進行注冊的位置在RegisterXLA.cpp這個文件中（codegen的結果），如下：

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框架內部的算子是完全一致的，只是修改了實現。而XLANativeFunctions::abs的實現可以在aten_xla_type.cpp中找到，如下所示：

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

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

總結

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

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就是根據xla_native_functions.yaml這個文件來生成算子的定義，然后再生成對應的RegisterXLA.cpp中的注冊代碼，這也跟PyTorch的codegen方式一致。

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

編輯：黃飛