Metadata-Version: 2.4
Name: custom-dl-optimizer
Version: 1.1.0
Summary: A profile-guided PyTorch inference optimizer using FX and Triton
Home-page: https://github.com/Devrajsinh-Jhala/Custom-DL-Optimizer
Author: Devrajsinh Jhala
License: MIT
Project-URL: Homepage, https://github.com/Devrajsinh-Jhala/Custom-DL-Optimizer
Project-URL: Source, https://github.com/Devrajsinh-Jhala/Custom-DL-Optimizer
Project-URL: Bug Tracker, https://github.com/Devrajsinh-Jhala/Custom-DL-Optimizer/issues
Project-URL: PyPI, https://pypi.org/project/custom-dl-optimizer/
Keywords: pytorch,triton,cuda,deep-learning,compiler,inference,optimization,fx
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Compilers
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1
Requires-Dist: triton>=2.1; platform_system != "Windows"
Provides-Extra: vision
Requires-Dist: torchvision>=0.16; extra == "vision"
Provides-Extra: dev
Requires-Dist: build>=1.0; extra == "dev"
Requires-Dist: twine>=4.0; extra == "dev"
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: torchvision>=0.16; extra == "dev"
Dynamic: author
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: home-page
Dynamic: keywords
Dynamic: license
Dynamic: license-file
Dynamic: project-url
Dynamic: provides-extra
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# Custom-DL-Optimizer

[![PyPI version](https://img.shields.io/pypi/v/custom-dl-optimizer.svg)](https://pypi.org/project/custom-dl-optimizer/)
[![Python](https://img.shields.io/pypi/pyversions/custom-dl-optimizer.svg)](https://pypi.org/project/custom-dl-optimizer/)
[![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](LICENSE)

Custom-DL-Optimizer is a lightweight, research-oriented PyTorch inference optimizer for NVIDIA GPUs. It applies profiling, PyTorch FX graph rewriting, custom Triton operator replacement, channels-last memory layout conversion, and automatic mixed precision to improve inference latency while preserving the normal PyTorch module interface.

The package is designed for ML systems experiments, compiler-oriented coursework, and prototype deployment studies where the model should remain a PyTorch `nn.Module` rather than being exported to an external runtime.

## Highlights

- PyTorch-native API: optimize an existing `torch.nn.Module` in one call.
- FX graph surgery: traces model graphs and rewrites supported operators.
- Triton kernel injection: replaces supported activations with custom GPU kernels when safe.
- Hardware-aware execution: supports channels-last memory layout and FP16 autocast.
- Safe fallbacks: CPU tensors and unsupported layouts fall back to PyTorch operators.
- Research notebook: includes a Colab benchmark workflow for reproducible tables and figures.

## Installation

```bash
pip install custom-dl-optimizer
```

For GPU acceleration, use a CUDA-enabled PyTorch build and a Linux/Colab environment with NVIDIA GPU support.

```bash
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
```

## Quick Start

```python
import torch
from torchvision.models import resnet50
from custom_dl_optimizer import AutoOptimizer

device = "cuda" if torch.cuda.is_available() else "cpu"
model = resnet50(weights=None).to(device).eval()
dummy_input = torch.randn(32, 3, 224, 224, device=device)

optimizer = AutoOptimizer(model, device=device)
optimized_model = optimizer.optimize(dummy_input)

with torch.inference_mode():
    baseline_output = model(dummy_input)
    optimized_output = optimized_model(
        dummy_input.to(memory_format=torch.channels_last)
        if dummy_input.is_cuda
        else dummy_input
    )

print(torch.allclose(baseline_output.float(), optimized_output.float(), rtol=8e-2, atol=8e-2))
```

## API

### `AutoOptimizer(model, device="cuda")`

Creates an optimizer wrapper around a PyTorch model.

Parameters:

- `model`: a `torch.nn.Module`.
- `device`: target device. Use `"cuda"` for GPU optimization or `"cpu"` for fallback behavior.

### `optimize(dummy_input)`

Runs the optimization pipeline and returns an optimized `nn.Module`.

Pipeline stages:

1. Profile the model with `torch.profiler`.
2. Trace the model into a PyTorch FX graph.
3. Replace supported modules, currently `nn.ReLU`, with optimized alternatives.
4. Convert CUDA models to channels-last memory format for convolution-heavy workloads.
5. Wrap CUDA inference in FP16 automatic mixed precision.

## Architecture

| Component | File | Responsibility |
| --- | --- | --- |
| Orchestrator | `custom_dl_optimizer/core/engine.py` | User-facing optimization pipeline |
| Profiler | `custom_dl_optimizer/core/profiler.py` | Runtime operator profiling and bottleneck reporting |
| Graph Surgeon | `custom_dl_optimizer/core/graph_surgeon.py` | FX tracing and module replacement |
| Triton Kernels | `custom_dl_optimizer/core/triton_kernels.py` | Custom GPU kernels and PyTorch fallbacks |

## Benchmark Snapshot

The repository includes `Custom_DL_Optimizer_Research_Colab.ipynb`, a Colab benchmark notebook that generates CSV, LaTeX tables, and publication-style graphs.

Example NVIDIA Tesla T4 results from the research notebook:

| Model | Batch | Eager FP32 (ms) | AMP/NHWC (ms) | Custom-DL (ms) | Speedup vs Eager | Speedup vs AMP |
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
| ResNet-50 | 128 | 430.317 | 161.772 | 109.243 | 3.94x | 1.48x |
| MobileNet-V2 | 128 | 124.204 | 62.109 | 35.989 | 3.45x | 1.73x |
| VGG-16 | 128 | 751.234 | 315.222 | 255.445 | 2.94x | 1.23x |
| EfficientNet-B0 | 128 | 171.223 | 80.883 | 45.675 | 3.75x | 1.77x |
| DenseNet-121 | 128 | 414.769 | 227.366 | 132.560 | 3.13x | 1.72x |

These results should be treated as hardware- and software-version-dependent measurements. For papers or reports, rerun the notebook on the target GPU and report the generated CSV values.

## Reproducible Benchmarking

1. Open `Custom_DL_Optimizer_Research_Colab.ipynb` in Google Colab.
2. Select `Runtime > Change runtime type > GPU`.
3. Restart the runtime.
4. Run all cells.
5. Use the generated files from:

```text
/content/custom_dl_optimizer_research_outputs
```

The notebook exports:

- `custom_dl_optimizer_research_results.csv`
- `custom_dl_optimizer_research_results.json`
- `paper_table_results.tex`
- `figure_1_latency_comparison.png`
- `figure_2_speedup.png`
- `figure_3_pass_coverage.png`
- `figure_4_output_parity.png`
- `generated_abstract_results.txt`

## Scope and Limitations

Custom-DL-Optimizer is a research prototype, not a replacement for TensorRT, Torch-TensorRT, TVM, XLA, or TorchInductor. It is most useful as a PyTorch-native experimentation layer for studying compiler passes and GPU-aware inference transformations.

Current limitations:

- The package-level optimizer currently focuses on CNN-style graphs and `nn.ReLU` replacement.
- Performance depends on GPU architecture, PyTorch version, CUDA version, tensor shapes, and batch size.
- User-defined Triton kernels may interact differently with `torch.compile` across PyTorch/Triton versions.
- Output parity should be checked with FP16-appropriate tolerances when AMP is enabled.

## Roadmap

- Add explicit TorchInductor and CUDA Graph integration to the package API.
- Add Conv-BatchNorm folding as a first-class graph pass.
- Add transformer-oriented LayerNorm/GELU and MLP block optimizations.
- Add post-training INT8 quantization experiments.
- Add CI tests for graph rewriting and output parity.
- Publish benchmark reports for T4, L4, A10, and A100 GPUs.

## Development

```bash
git clone https://github.com/Devrajsinh-Jhala/Custom-DL-Optimizer.git
cd Custom-DL-Optimizer
python -m pip install -e ".[dev]"
pytest
```

Build the package locally:

```bash
python -m build
python -m twine check dist/*
```

Upload to PyPI after checking the package description:

```bash
python -m twine upload dist/*
```

## Citation

If you use this project in a paper, report the exact package version, GPU, CUDA version, PyTorch version, batch size, warmup count, benchmark iterations, and output-parity tolerance.

Suggested citation text:

```text
Custom-DL-Optimizer: A profile-guided PyTorch micro-compiler for hardware-aware NVIDIA GPU inference optimization.
```

## License

This project is released under the MIT License. See [LICENSE](LICENSE).
