一、背景說明

1.1 PyTorch擴展機制

• PrivateUse1是PyTorch為第三方設備擴展設計的保留設備類型，允許開發者添加新硬件支持
• 當算子在當前設備（如NPU）未實現時，PyTorch會自動回退（fallback）到CPU執行
• 本文以native_batch_norm算子為例，演示如何為NPU設備添加自定義實現

1.2 核心概念

• ATen：PyTorch的核心張量運算庫，提供超過2000個基礎算子
• 內存格式：描述張量在內存中的排布方式，如NCHW（批處理x通道x高度x寬度）
• 自動微分：PyTorch通過記錄計算圖實現反向傳播，需要同時實現前向和反向算子

二、實現步驟詳解

2.1 實現前向、反向傳播算子

cat > native_batch_norm_npu.cpp <<-'EOF'// 包含必要的頭文件
#include <torch/library.h>      // 算子注冊相關
#include <ATen/EmptyTensor.h>   // 空張量創建
#include <ATen/Device.h>
#include <ATen/Utils.h>
#include <ATen/native/Resize.h>
#include <c10/core/DeviceType.h> // 設備類型定義std::tuple<at::Tensor, at::Tensor, at::Tensor> native_batch_norm_npu(const at::Tensor& input,  // 輸入張量const c10::optional<at::Tensor>& weight, // 可選的gamma參數（縮放）const c10::optional<at::Tensor>& bias,   // 可選的beta參數（偏移）const c10::optional<at::Tensor>& running_mean, // 訓練時累計的均值const c10::optional<at::Tensor>& running_var,  // 訓練時累計的方差bool training, // 是否處于訓練模式double momentum, // 動量參數，用于更新running統計量double eps)      // 數值穩定系數
{// 創建臨時張量占位（實際實現需計算真實統計量）at::Tensor output = at::empty_like(input);at::Tensor dummy_mean = at::empty