"""
custom_ops.py
=============
Loads the C++ / CUDA shared library and sets up all custom PyTorch operations
used by the IdentityModel:

  1. torch.classes.my_ops.IdentityConvClass  (registered by the shared library)
     - A fake/abstract version is registered here so that torch.export can
       trace through module attributes that hold an instance of this class.

  2. my_ops::identity_conv_op  (schema + CPU + CUDA registered by the shared library)
     - register_fake: abstract implementation for torch.export / FakeTensor.
"""

import os

import torch
import torch.library











_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)











from torch._library.fake_class_registry import register_fake_class


@register_fake_class("my_ops::IdentityConvClass")
class FakeIdentityConvClass:
    """Abstract counterpart of IdentityConvClass used during 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_











@torch.library.register_fake("my_ops::identity_conv_op")
def _identity_conv_op_fake(x: torch.Tensor) -> torch.Tensor:
    """Abstract implementation used by torch.export / FakeTensor tracing."""
    return torch.empty_like(x)



identity_conv_op = torch.ops.my_ops.identity_conv_op

"""
model.py
========
Defines the four-layer IdentityModel used in the AOTInductor demo.

Layer layout
------------
  layer1 : IdentityConv           - native PyTorch operators
  layer2 : IdentityConvCustomClass - torch.classes  C++/CUDA custom class
  layer3 : IdentityConvCustomOp   - torch.library.custom_op  C++/CUDA op
  layer4 : IdentityConv           - native PyTorch operators

Every layer is an identity transformation, so model(x) == x for any input x.
"""

import torch
import torch.nn as nn



from custom_ops import identity_conv_op  





class IdentityConv(nn.Module):
    """Identity convolution implemented with native PyTorch operators.

    Uses a depthwise Conv2d with kernel_size=1 and weight=1.0, which is
    equivalent to a no-op (output == input).  This layer is compatible with
    torch.export and AOTInductor out of the box.
    """

    def __init__(self, channels: int) -> None:
        super().__init__()
        self.conv = nn.Conv2d(
            in_channels=channels,
            out_channels=channels,
            kernel_size=(1, 1),
            stride=(1, 1),
            padding=(0, 0),
            dilation=(1, 1),
            groups=channels,
            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):
    """Identity convolution backed by a torch.classes C++/CUDA custom class.

    At runtime the forward call is dispatched to the CUDA kernel registered
    inside IdentityConvClass (csrc/identity_conv.cpp + .cu).

    For torch.export compatibility a FakeIdentityConvClass is registered in
    custom_ops.py via @register_fake_class so that symbolic tracing works.
    """

    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):
    """Identity convolution backed by a torch.library.custom_op C++/CUDA op.

    The op (my_ops::identity_conv_op) is defined in custom_ops.py with:
      • a register_fake  implementation for torch.export tracing
      • a register_kernel("cuda") implementation that calls the CUDA kernel
    """

    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):
    """Four-layer identity model for AOTInductor demo."""

    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 an IdentityModel in eval mode on the default CUDA device."""
    return IdentityModel(channels=channels).cuda().eval()

"""
export_compile.py
=================
Exports the IdentityModel with torch.export and compiles it with
torch._inductor.aoti_compile_and_package.

The resulting package (model.pt2) is written to the artifacts/ directory
and can be loaded by both run_inference.py (Python) and the C++ inference
binary.

Usage (run from the python/ directory):
    python export_compile.py
"""

import os
import sys


sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import torch


import custom_ops  
from model import create_model




CHANNELS = 3
BATCH_SIZE = 1
HEIGHT = 224
WIDTH = 224


PACKAGE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..",
                            "artifacts", "model.pt2")


def main() -> None:
    print("=" * 64)
    print("AOTInductor  -  Export & Compile")
    print("=" * 64)

    
    
    
    print(f"\n[1/4] Creating IdentityModel  (channels={CHANNELS}) ...")
    model = create_model(channels=CHANNELS)

    x = torch.randn(BATCH_SIZE,
                    CHANNELS,
                    HEIGHT,
                    WIDTH,
                    device="cuda",
                    dtype=torch.float32)

    with torch.no_grad():
        out = model(x)

    assert torch.equal(
        x, out), (f"Eager pre-export check FAILED  "
                  f"(max diff = {(x - out).abs().max().item():.2e})")
    print("  Eager verification  PASSED  (bitwise identical)")

    
    
    
    print("\n[2/4] Exporting model with torch.export.export() ...")
    with torch.no_grad():
        exported_program = torch.export.export(model, (x, ))
    print("  Export  DONE")
    print(f"\n  Exported graph:\n{exported_program.graph}")

    
    
    
    print(
        "\n[3/4] Compiling with torch._inductor.aoti_compile_and_package ...")
    package_path = torch._inductor.aoti_compile_and_package(
        exported_program,
        package_path=PACKAGE_PATH,
    )
    print(f"  Compilation  DONE")
    print(f"  Package saved to: {os.path.abspath(package_path)}")

    
    
    
    print(
        "\n[4/4] Quick sanity check: loading package and running inference ..."
    )
    compiled_model = torch._inductor.aoti_load_package(package_path)
    with torch.no_grad():
        out_compiled = compiled_model(x)

    
    if isinstance(out_compiled, (list, tuple)):
        out_compiled = out_compiled[0]

    assert torch.equal(x, out_compiled), (
        f"Compiled model sanity check FAILED  "
        f"(max diff = {(x - out_compiled).abs().max().item():.2e})")
    print("  Compiled model verification  PASSED  (bitwise identical)")

    print("\n" + "=" * 64)
    print(f"SUCCESS!  Package: {os.path.abspath(package_path)}")
    print("=" * 64)


if __name__ == "__main__":
    main()

"""
run_inference.py
================
Loads the AOTInductor-compiled IdentityModel package (model.pt2) and runs
inference to verify correctness.

The output of the identity model must equal the input within a tight
floating-point tolerance.

Usage (run from the python/ directory after export_compile.py):
    python run_inference.py [MODEL_PATH [OP_LIB_PATH]]

Arguments:
    MODEL_PATH   Path to the compiled model package (.pt2).
                 Defaults to ../artifacts/model.pt2 relative to this script.
    OP_LIB_PATH  Path to the custom-op shared library (.so).
                 When provided the library path is forwarded to custom_ops.py
                 via the IDENTITY_CONV_OPS_LIB environment variable so that
                 torch.ops.load_library uses that file instead of the default
                 ../ext/libidentity_conv_ops.so.
"""

import os
import sys


sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import torch
import torch._inductor.codecache  




_DEFAULT_PACKAGE_PATH = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "..", "artifacts", "model.pt2")

PACKAGE_PATH = sys.argv[1] if len(sys.argv) > 1 else _DEFAULT_PACKAGE_PATH
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"] = os.path.abspath(OP_LIB_PATH)



import custom_ops  




CHANNELS = 3
BATCH_SIZE = 1
HEIGHT = 224
WIDTH = 224


def main() -> None:
    print("=" * 64)
    print("AOTInductor  -  Python Inference")
    print("=" * 64)

    
    
    
    pkg = os.path.abspath(PACKAGE_PATH)
    if OP_LIB_PATH is not None:
        print(f"  Op library  : {os.path.abspath(OP_LIB_PATH)}")
    print(f"\n[1/3] Loading compiled model from:\n  {pkg}")
    compiled_model = torch._inductor.aoti_load_package(pkg)
    print("  Model loaded successfully.")

    
    
    
    x = torch.randn(BATCH_SIZE,
                    CHANNELS,
                    HEIGHT,
                    WIDTH,
                    device="cuda",
                    dtype=torch.float32)
    print(f"\n[2/3] Input  shape={list(x.shape)}  dtype={x.dtype}  "
          f"device={x.device}")

    
    
    
    print("\n[3/3] Running inference ...")
    with torch.no_grad():
        output = compiled_model(x)

    
    if isinstance(output, (list, tuple)):
        output = output[0]

    print(f"  Output shape={list(output.shape)}  dtype={output.dtype}")

    if torch.equal(x, output):
        print("\n  Verification  PASSED  (bitwise identical)")
    else:
        max_diff = (x - output).abs().max().item()
        print(f"\n  Verification  FAILED  (max diff = {max_diff})"
              f" — expected bitwise identical output")
        sys.exit(1)

    print("\n" + "=" * 64)
    print("SUCCESS!  AOTInductor Python inference verified.")
    print("=" * 64)


if __name__ == "__main__":
    main()