---
language:
  - en
library_name: hls4ml
datasets:
  - lithobench
tags:
  - pytorch
  - hls4ml
  - fpga
  - neural-network
  - quantization
  - xilinx
  - mask-optimization
  - lithography
  - inverse-lithography
license: gpl
---

# Penumbra UNet: FPGA-Accelerated Mask Optimization

A compressed U-Net neural network for on-chip FPGA acceleration of Inverse Lithography Technology (ILT) mask optimization, targeting the Xilinx VU47P (AWS F2).

## Overview

Penumbra UNet compresses a full-size teacher network by 64× (7.8M → 122K parameters) to fit entirely in on-chip BRAM, enabling a fully on-chip dataflow that eliminates external DRAM access.

## Architecture

### Network Structure

U-Net encoder-decoder with extreme parameter compression:

**Encoder:**
- Conv 1→8 channels, 64×64 + MaxPool → 32×32
- Conv 8→16 channels, 32×32 + MaxPool → 16×16
- Conv 16→32 channels, 16×16 + MaxPool → 8×8

**Bottleneck:**
- Conv 32→64 channels, 8×8

**Decoder:**
- Upsample + skip concatenation + Conv 96→32 channels, 16×16
- Upsample + skip concatenation + Conv 48→16 channels, 32×32
- Upsample + skip concatenation + Conv 24→8 channels, 64×64

**Output:**
- Conv 1×1 + Sigmoid → 64×64

**Compression summary:**
| Metric | Full model | Penumbra UNet |
|--------|-----------|----------|
| Parameters | 7.8M | 122K |
| Input tile | 512×512 | 64×64 |
| Max channels | 512 | 64 |

### Tiling & Reassembly

Input 512×512 masks are decomposed into 16×16 grid of 64×64 tiles (256 total):
- **Overlap**: 16-pixel reflection padding for boundary handling
- **Usable core**: 32×32 center pixels per tile
- **Batch processing**: 256 tiles → 4 sequential batches of 64

Reassembly uses only differentiable operations (slice, reshape, permute) to enable end-to-end gradient flow:
```
(256, 1, 64, 64)  [all tiles]
    ↓ center-crop
(256, 1, 32, 32)  [usable cores]
    ↓ reshape + permute
(1, 1, 512, 512)  [full mask]
```

## Training

### Phase 1: Knowledge Distillation
- **Epochs**: 16
- **Input**: 64×64 crops
- **Loss**: α-blended (α decays 0.7→0)
  ```
  L = α·MSE(student, teacher) + (1-α)·MSE(student, ground_truth)
  ```
- **Optimizer**: Adam (lr=1e-3), cosine-annealing schedule
- **Teacher**: Frozen full-size NeuralILT model

### Phase 2: Physics-Informed Fine-Tuning
- **Epochs**: 4
- **Pipeline**: Full tiled forward pass through differentiable lithography simulator
- **Loss**: Print fidelity + process variation
  ```
  L = MSE(P_nom, target) + MSE(P_max, P_min)
  ```
- **Optimizer**: Adam (lr=1e-4), StepLR (γ=0.1 at epoch 2)
- **Gradients**: Flow through tiled reassembly to all network weights

## Code Organization

```
hls4ml_penumbra/
├── firmware/           # Generated HLS C++ project
│   ├── myproject.cpp   # Top-level module
│   ├── myproject.h     # Interface & config
│   ├── weights/        # Quantized weights
│   ├── ap_types/       # Xilinx AP types (ap_fixed, ap_int)
│   └── utils/          # HLS utilities
├── myproject_prj/      # Vivado HLS project
│   └── solution1/
│       └── impl/       # Implementation artifacts
├── logs/               # Build logs
└── [HLS build outputs]
```

---

**Author**: Roberto Treviño Cervantes