目錄
快捷方式

後端方言

概覽

後端方言是 edge dialect 的一種特殊變體，因為它在後端特定的圖轉換後，包含後端特定的節點和元資料。後端方言是一個可選階段，只有當我們想在圖中引入後端感知（backend-awareness）時才需要。更具體地說，後端方言中的圖可能包含僅對目標後端有意義的 operator 或委託的下沉模組（參見 delegate 文件）。一個用例是，如果我們想將 operator 融合到一個 operator 中，例如，將連續的 addmm + relu 融合到單個 operator addmm_relu 中，我們可以在此處進行。

本文件描述瞭如何引入後端特定的 operator。

自定義 operator 和後端特定 operator 的區別：自定義 operator 會出現在 eager 模式、ATen 方言和 edge 方言中，而後端特定 operator 僅由發生在 edge 方言之後的 pass 引入。

何時使用

此方言允許引入不符合 canonical ATen operator set 中定義的 schema 且不出現在上述任何方言（ATen 方言和 edge 方言）中的 operator。如果您的用例滿足以下一個或多個標準，請考慮使用後端 operator

  • 您的後端提供了一個庫，該庫優化了等同於子圖的某個 operator。例如，linear_relu（等同於 linear + relu），可以在特定後端上執行得更快。

  • 將圖模組下沉到後端後，需要重新追蹤（retrace）它。當我們重新追蹤時，後端 operator 可以轉換回原始子圖（在 ATen 方言中），而普通的自定義 operator 不會處理這種情況。

  • 您的後端特定 operator 沒有通用的 CPU kernel，只有針對特定後端的 kernel。使用後端 operator 可以透過使用原始子圖作為預設 kernel 並保持圖模組可執行來解決此問題。

  • 或者，如果您擔心這可能過於複雜，並且只想要更輕量級且僅在編譯器階段需要 Python 程式碼的東西，則可以使用 delegate。

API

對於 operator/子圖替換，通用流程是

  1. 註冊一個與子圖具有相同輸入和輸出的 operator。這個 operator 不需要具有目標特定的實現（在編譯階段也不需要），但它需要提供與子圖相同的結果。

  2. 建立一個 pattern，使編譯器能夠找到該子圖並將其替換為替代項。

  3. 編寫一個 pass，用新的 operator 替換子圖。

為了方便此過程，我們提供了一個 API 來幫助 ExecuTorch 使用者減少執行這些步驟的工作量。

Pass 基礎設施入口點

要將 edge operator 下沉到 backend operator，pass 將執行 pattern 匹配，識別圖中的目標 edge operator，然後用等效的 backend operator 替換它們。有兩種 API 可用於註冊此類 pass

  • transform()。這是 ExportProgram 上的一個 API，允許使用者提供自定義 pass。請注意，此 API 不受任何 validator 的保護，因此無法保證程式的正確性（soundness）。

  • ExecutorchBackendConfig.passes。如果在此處新增，該 pass 將成為從 backend 方言到 ExecutorchProgram 的下沉過程的一部分。

示例：一個這樣的 pass 是 QuantFusion。這個 pass 接收一個“規範量化 pattern”，即“dequant - some_op - quant”，並將這個 pattern 融合到一個後端特定的 operator 中，例如 quantized_decomposed::some_op。另一個更簡單的示例在此處：here，我們將 sym_size operator 替換為 ExecuTorch 可以理解的 operator。

Pattern 繫結 Decorator

我們提供了一個 decorator bind_pattern_to_op 來幫助使用者輕鬆地將其後端 operator 註冊到 EXIR 中。此 decorator 接受

  • 一個 torch.Library 物件，它指示此後端 operator 屬於哪個 library 或 namespace。

  • 一個名稱或 schema。如果已經在 torch.Library 物件中定義了後端 operator 的 schema，則只需要提供一個名稱。否則，如果傳入 schema 字串，我們可以註冊該 schema。

此 decorator 應新增到我們在 edge 方言中嘗試匹配（然後下沉到此後端 operator）的 pattern 上。透過這種方式，我們將此 pattern 註冊為該後端 operator 的 CompositeImplicitAutograd kernel。

然後可以從 pass 中訪問/使用該 operator。CompositeImplicitAutograd kernel 確保

  1. 使用者無需編寫 (CPU) 可執行的 kernel。

  2. 確保 ExportProgram 的可重新追蹤性（retrace-ability）。一旦重新追蹤，後端 operator 將被分解回 pattern 中使用的 ATen operator。

示例

讓我們假設一個包含 add 和 relu 兩個 operator 的簡單程式

def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    z = x + y
    return torch.ops.aten.relu.default(z)

下沉到 edge 方言後，它變為

graph():
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %arg1_1 : [num_users=1] = placeholder[target=arg1_1]
    %aten_add_tensor : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.add.Tensor](args = (%arg0_1, %arg1_1), kwargs = {})
    %aten_relu_default : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.relu.default](args = (%aten_add_tensor,), kwargs = {})
    return (aten_relu_default,)

現在我想編寫一個 pass，將 add 和 relu 合併到 add_relu 中，第一步是編寫一個 pattern

# In the pattern, we can use edge ops and ATen ops interchangably
def pattern(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    z = torch.ops.aten.add.Tensor(x, y)
    out = torch.ops.aten.relu.default(z)
    return out

然後我們需要從融合後的 operator namespace 建立一個 operator library，然後將 decorator 應用到我們的 pattern 上

lib = Library("foo_namespace", "DEF")

@bind_pattern_to_op(lib, "add_relu(Tensor self, Tensor other) -> Tensor")
def pattern(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        z = torch.ops.aten.add.Tensor(x, y)
        out = torch.ops.aten.relu.default(z)
        return out

透過這種方式，我們將 pattern 註冊為 add_relu 的一個 kernel，並已準備好在 pass 中使用。一個簡單的 pass 如下所示

class AddReluFusionPass(ExportPass):
    def call(self, graph_module: GraphModule) -> PassResult:
        # decorator registers this pattern as a CompositeExplicitAutograd kernel, since there's no kernel registered before.
        @bind_pattern_to_op(lib, "add_relu")
        def pattern(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
            z = torch.ops.aten.add.Tensor(x, y)
            out = torch.ops.aten.relu.default(z)
            return out

        def replacement(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
            return torch.ops.foo_namespace.add_relu.default(x, y)

        subgraph_rewriter.replace_pattern(
            graph_module,
            _trace_and_lower_to_edge_ops(pattern),
            _trace_and_lower_to_edge_ops(replacement),
        )
        return PassResult(graph_module, True)

結果圖如下所示

graph():
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %arg1_1 : [num_users=1] = placeholder[target=arg1_1]
    %foo_namespace_add_relu_default : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.foo_namespace.add_relu.default](args = (%arg0_1, %arg1_1), kwargs = {})
    return (foo_namespace_add_relu_default,)

Operator Set

以下是當前使用 bind_pattern_to_op API 的後端 operator。

  • executorch_prims::add.int(SymInt a, SymInt b) -> SymInt

    • pattern: builtin.add

    • backend: executor

  • executorch_prims::mul.int(SymInt a, SymInt b) -> SymInt

    • pattern: builtin.mul

    • backend: executor

  • executorch_prims::sub.int(SymInt a, SymInt b) -> SymInt

    • pattern: builtin.sub

    • backend: executor

  • executorch_prims::floordiv.int(SymInt a, SymInt b) -> SymInt

    • pattern: builtin.floordiv

    • backend: executor

  • executorch_prims::truediv.int(Scalar a, Scalar b) -> Scalar

    • pattern: builtin.div

    • backend: executor

  • executorch_prims::sym_float.Scalar(Scalar a) -> Scalar

    • pattern: builtin.float

    • backend: executor

  • executorch_prims::gt.int(SymInt a, SymInt b) -> bool

    • pattern: builtin.gt

    • backend: executor

  • executorch_prims::lt.int(SymInt a, SymInt b) -> bool

    • pattern: builtin.lt

    • backend: executor

  • executorch_prims::ge.int(SymInt a, SymInt b) -> bool

    • pattern: builtin.ge

    • backend: executor

  • executorch_prims::le.int(SymInt a, SymInt b) -> bool

    • pattern: builtin.le

    • backend: executor

  • executorch_prims::eq.int(SymInt a, SymInt b) -> bool

    • pattern: builtin.eq

    • backend: executor

  • executorch_prims::mod.Scalar(SymInt a, SymInt b) -> SymInt

    • pattern: builtin.divmod

    • backend: executor

  • executorch_prims::neg.Scalar(Scalar a) -> Scalar

    • pattern: operator.ne

    • backend: executor

  • quantized_decomposed::embedding_byte(Tensor weight, Tensor weight_scales, Tensor weight_zero_points, int weight_quant_min, int weight_quant_max, Tensor indices) -> Tensor

    • pattern: source

    • backend: quantization

  • quantized_decomposed::add(Tensor a, float a_scale, int a_zero_point, int a_quant_min, int a_quant_max, Tensor b, float b_scale, int b_zero_point, int b_quant_min, int b_quant_max, float out_scale, int out_zero_point, int out_quant_min, int out_quant_max) -> Tensor qc

    • pattern: source

    • backend: quantization

  • quantized_decomposed::add.scalar(Tensor qa, float a_scale, int a_zero_point, int a_quant_min, int a_quant_max, ScalarType a_dtype, Scalar b, float out_scale, int out_zero_point, int out_quant_min, int out_quant_max, ScalarType out_dtype) -> Tensor

    • pattern: source

    • backend: quantization

  • quantized_decomposed::add_relu(Tensor a, float a_scale, int a_zero_point, int a_quant_min, int a_quant_max, Tensor b, float b_scale, int b_zero_point, int b_quant_min, int b_quant_max, float out_scale, int out_zero_point, int out_quant_min, int out_quant_max) -> Tensor qc

    • pattern: source

    • backend: quantization

文件

訪問 PyTorch 的全面開發者文件

檢視文件

教程

獲取面向初學者和高階開發者的深度教程

檢視教程

資源

查詢開發資源並獲得問題解答

檢視資源

為了分析流量和最佳化您的體驗，我們在本網站上使用 cookie。點選或導航即表示您同意我們使用 cookie。作為本網站的當前維護者，Facebook 的 Cookies Policy 適用。瞭解更多資訊，包括關於可用控制元件的資訊：Cookies Policy。