PyTorch カスタムオペレーションと AOTInductor で動作させる方法

2026 年 5 月 10 日に公開された記事で、C++ と CUDA で実装した PyTorch カスタムオペレーション（関数）およびカスタムクラスの具体的な実装例を踏まえ、AOTInductor コンパイルと推論プログラムでの利用方法を紹介する。

PyTorch のカスタムオペレーションとは、

C++

CUDA

概要：恒等畳み込みの実装例

本記事は、単純な**「恒等演算（input == output）」による畳み込み**を実際に作成し、以下の手順で扱うことを示します。

• C++ と CUDA で実装するカスタムオペレーションの構築
• 登録・ロード処理を通じた PyTorch モデルへの統合

• torch.compile

  および

  torch.export

  との完全な互換性確保
• Python および純粋な C++ 推論プログラムでの実行

1. PyTorch カスタム関数の実装（C++ / CUDA）

カスタム関数はステートレスであり、パラメータを持たせません。CPU および CUDA の両方の実装を提供し、

TORCH_LIBRARY_IMPL

実装コード (

custom_ops.cpp

)

以下のコードは、GPU/ CPU で動作する「恒等演算」のロジックを定義しています。

// custom_ops.cpp
#include <torch/extension.h>
#include <cuda_runtime.h>

namespace {

// CUDA カーネル：並列処理で要素ごとのコピーを実行
template<typename scalar_t>
__global__ void identity_kernel(const scalar_t* input, scalar_t* output,
                                const int64_t* shape_dev, const int64_t* strides_dev,
                                int ndim, int numel) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < numel) {
        // 恒等演算：入力値をそのまま出力にコピー
        output[idx] = input[idx];
    }
}

// CPU 用実装：テンソルをクローン
torch::Tensor identity_conv_cpu_impl(const torch::Tensor& input){
    TORCH_CHECK(!input.is_cuda(), "identity_conv_cpu_impl: input must be a CPU tensor");
    return input.clone();
}

// CUDA 用実装
torch::Tensor identity_conv_cuda_impl(const torch::Tensor& input){
    TORCH_CHECK(input.is_cuda(), "identity_conv_cuda_impl: input must be a CUDA tensor");
    
    auto output = torch::empty_like(input);
    const int64_t numel = input.numel();
    if (numel == 0) return output;
        
    const int ndim = input.dim();
    // 形状とストライドをデバイス上へコピー（カーネル引数用）
    const auto opts = torch::TensorOptions().dtype(torch::kInt64).device(input.device());
    const auto shape_dev = torch::tensor(
        std::vector<int64_t>(input.sizes().begin(), input.sizes().end()), opts);
    const auto strides_dev = torch::tensor(
        std::vector<int64_t>(input.strides().begin(), input.strides().end()), opts);
    
    constexpr int kThreads = 256;
    const int blocks = static_cast<int>((numel + kThreads - 1) / kThreads);
    
    AT_DISPATCH_FLOATING_TYPES_AND2(
        at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(),
        "identity_conv_cuda_impl",
        [&]() {
            identity_kernel<scalar_t><<<blocks, kThreads>>>(
                input.data_ptr<scalar_t>(), output.data_ptr<scalar_t>(),
                shape_dev.data_ptr<int64_t>(), strides_dev.data_ptr<int64_t>(),
                ndim, numel);
        });
    
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}

// 登録処理：デバイス種別ごとにマクロを使用
TORCH_LIBRARY_IMPL(my_ops, CUDA, m){
    m.impl("identity_conv_op", identity_conv_cuda_impl);
}

TORCH_LIBRARY_IMPL(my_ops, CPU, m){
    m.impl("identity_conv_op", identity_conv_cpu_impl);
}

} // namespace

2. PyTorch カスタムクラスの実装（C++）

パラメータや

forward()

アプローチ

• C++ で

  torch::CustomClassHolder

  を継承してクラスを定義

• TORCH_LIBRARY

  マクロを使用して登録
• Python からの呼び出し（PyTorch グラフ）に対応させる

実装コード (

custom_class.cpp

)

// custom_class.cpp
#include <torch/extension.h>

struct IdentityConvClass : torch::CustomClassHolder{
    int64_t channels_;
    // コンストラクタ：パラメータを保持
    explicit IdentityConvClass(int64_t channels) : channels_(channels) {}
    
    // フォワード処理：デバイス検知付き
    torch::Tensor forward(const torch::Tensor& x) {
        return x.is_cuda() ? identity_conv_cuda_impl(x)
                           : identity_conv_cpu_impl(x);
    }
    
    int64_t get_channels() const { return channels_; }
};

// custom_class_registration.cpp
#include <torch/extension.h>
#include <c10/core/TensorImpl.h>

TORCH_LIBRARY(my_ops, m){
            // クラスの登録
            m.class_<IdentityConvClass>("IdentityConvClass")
        .def(torch::init<int64_t>()) // コンストラクタバインディング
        .def("forward", &IdentityConvClass::forward) // メソッドバインディング
        .def("get_channels", &IdentityConvClass::get_channels)
        .def("__obj_flatten__", [](const c10::intrusive_ptr<IdentityConvClass>& self){
                 return std::make_tuple(std::make_tuple("channels", self->channels_));
             })
        // Pickle シリアライザの登録（保存・復元用）
        .def_pickle(
            [](const c10::intrusive_ptr<IdentityConvClass>& self) -> int64_t { return self->channels_; },
            [](int64_t channels) -> c10::intrusive_ptr<IdentityConvClass> 
            { return c10::make_intrusive<IdentityConvClass>(channels); });
        
        // シンプルな関数のバインディングも定義可能
        m.def("identity_conv_op(Tensor x) -> Tensor");
}

3. Python 環境でのセットアップとファイク登録

C++ 側をビルドし、共有ライブラリ（

libidentity_conv_ops.so

torch.export

torch.compile

共有ライブラリのロードとファイク実装 (

custom_ops.py

)

# custom_ops.py
import os
import torch
from torch.library import default_lib

# 共有ライブラリのパスを設定
_default_lib = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "..", "ext",
    "libidentity_conv_ops.so")
_lib_path = os.path.abspath(
    os.environ.get("IDENTITY_CONV_OPS_LIB", _default_lib))

# ライブラリの読み込み
torch.ops.load_library(_lib_path)

# 1. カスタムクラスのファイク登録
# torch.export がこのクラスをグラフとして扱うために必須
from torch._library.fake_class_registry import register_fake_class

@register_fake_class("my_ops::IdentityConvClass")
class FakeIdentityConvClass:
    """torch.export 中の抽象的な対応物"""
    def __init__(self, channels: int) -> None:
        self.channels_ = channels
    
    def __obj_flatten__(self):
        return (("channels", self.channels_), )
    
    @classmethod
    def __obj_unflatten__(cls, flat):
        return cls(dict(flat)["channels"])
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 抽象実装：出力サイズは入力と同じ
        return torch.empty_like(x)
    
    def get_channels(self) -> int:
        return self.channels_

# 2. カスタム関数のファイク登録
@torch.library.register_fake("my_ops::identity_conv_op")
def _identity_conv_op_fake(x: torch.Tensor) -> torch.Tensor:
    """torch.export/FakeTensor 追跡用抽象実装"""
    return torch.empty_like(x)

identity_conv_op = torch.ops.my_ops.identity_conv_op

4. PyTorch モデルの構築（

IdentityModel

）

C++/CUDA を介したカスタムクラス・関数と、ネイティブ PyTorch オペレーターのミックスでモデルを作成します。

モデル定義 (

model.py

)

# model.py
import torch
import torch.nn as nn
from custom_ops import identity_conv_op

class IdentityConv(nn.Module):
    """ネイトヴ PyTorch で実装（kernel_size=1）"""
    def __init__(self, channels: int) -> None:
        super().__init__()
        self.conv = nn.Conv2d(channels, channels, kernel_size=(1, 1), stride=(1, 1), 
                              padding=(0, 0), bias=False)
        self.conv.weight.data = torch.ones(channels, 1, 1, 1)
        self.conv.weight.requires_grad = False
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.conv(x)

class IdentityConvCustomClass(nn.Module):
    """torch.classes で登録されたカスタムクラスを使用"""
    def __init__(self, channels: int) -> None:
        super().__init__()
        self.obj = torch.classes.my_ops.IdentityConvClass(channels)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.obj.forward(x)

class IdentityConvCustomOp(nn.Module):
    """torch.library.custom_op で登録された関数を使用"""
    def __init__(self, channels: int) -> None:
        super().__init__()
        self.channels = channels
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return identity_conv_op(x)

class IdentityModel(nn.Module):
    """AOTInductor デモ用：4 レイヤ構成"""
    def __init__(self, channels: int) -> None:
        super().__init__()
        self.layer1 = IdentityConv(channels)        # ネイティブ
        self.layer2 = IdentityConvCustomClass(channels) # カスタムクラス
        self.layer3 = IdentityConvCustomOp(channels)  # カスタム関数
        self.layer4 = IdentityConv(channels)        # ネイティブ
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        return x

def create_model(channels: int = 3) -> IdentityModel:
    return IdentityModel(channels=channels).cuda().eval()

5.

torch.export

と AOTInductor コンパイル

PyTorch モデルをエクスポートし、AOTInductor でコンパイルしてパッケージ化（

.pt2

エクスポートとコンパイル (

export_compile.py

)

# export_compile.py
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import torch
from model import create_model

CHANNELS = 3
BATCH_SIZE = 1
HEIGHT = 224
WIDTH = 224
PACKAGE_PATH = "./artifacts/model.pt2"

def main() -> None:
    model = create_model(channels=CHANNELS)
    
    # [入力生成]
    x = torch.randn(BATCH_SIZE, CHANNELS, HEIGHT, WIDTH, device="cuda", dtype=torch.float32)
    
    # [動作確認：Eager モード]
    with torch.no_grad(): out = model(x)
    assert torch.equal(x, out), "Eager verification FAILED"

    # [エクスポート: torch.export]
    with torch.no_grad():
        exported_program = torch.export.export(model, (x,))
    
    print("Exported graph:")
    # 出力グラフのサマリー（確認用）
    # %call_torchbind : IdentityConvClass の呼び出しを表す
    # %identity_conv_op : Custom Op の呼び出しを表す

    # [コンパイル: AOTInductor]
    torch._inductor.aoti_compile_and_package(exported_program, package_path=PACKAGE_PATH)
    print("Compilation DONE")

if __name__ == "__main__": main()

注意: エクスポートされたグラフを確認すると、カスタムクラスは

torch.ops.higher_order.call_torchbind

、カスタム関数は

torch.ops.my_ops.identity_conv_op.default

として表記されます。

6. 推論の実行

コンパイルされたパッケージ（

.pt2

A. Python 推論 (

run_inference.py

)

# run_inference.py
import sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import torch
from torch._inductor import codecache

PACKAGE_PATH = sys.argv[1] if len(sys.argv) > 1 else "./artifacts/model.pt2"
OP_LIB_PATH = sys.argv[2] if len(sys.argv) > 2 else None

if OP_LIB_PATH is not None:
    os.environ["IDENTITY_CONV_OPS_LIB"] = OP_LIB_PATH

compiled_model = torch._inductor.aoti_load_package(PACKAGE_PATH)
x = torch.randn(1, 3, 224, 224, device="cuda", dtype=torch.float32)
with torch.no_grad(): output = compiled_model(x)
assert torch.equal(x, output[0]), "Bitwise mismatch!"

print("SUCCESS! Python inference verified.")

B. C++ 推論 (

run_inference.cpp

)

純粋な C++ 環境（pybind11 不要）でも動作します。

dlopen

// run_inference.cpp
#include <cstdlib>
#include <iostream>
#include <stdexcept>
#include <dlfcn.h> 
#include <torch/torch.h>
#include <torch/csrc/inductor/aoti_package/model_package_loader.h>

static constexpr int64_t kBatchSize = 1;
static constexpr int64_t kChannels = 3;
static constexpr int64_t kHeight = 224;
static constexpr int64_t kWidth = 224;

int main(int argc, char* argv[]){
    if (argc < 3) {
        std::cerr << "Usage: ./run_inference <model.pt2> <libidentity_conv_ops.so>" << std::endl;
        return EXIT_FAILURE;
    }

    void* lib_handle = dlopen(argv[2], RTLD_NOW | RTLD_GLOBAL);
    if (!lib_handle) { throw std::runtime_error(std::string("dlopen failed: ") + dlerror()); }

    torch::inductor::AOTIModelPackageLoader loader(argv[1]);
    auto runner = loader.get_runner();

    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA, 0);
    auto input = torch::randn({kBatchSize, kChannels, kHeight, kWidth}, options);

    std::vector<at::Tensor> inputs = {input};
    auto outputs = runner->run(inputs);
    const auto& output = outputs[0];

    bool passed = input.equal(output);
    
    // リソース解放
    dlclose(lib_handle);

    if (passed) std::cout << "SUCCESS! C++ inference verified." << std::endl;
    else { std::cerr << "Verification FAILED" << std::endl; return EXIT_FAILURE; }

    return EXIT_SUCCESS;
}