Accelerating CRONet on AMD Versal AIE-ML Engines

Abstract: Topology optimization is a computational method used to determine the optimal material distribution within a prescribed design domain, aiming to minimize structural weight while satisfying load and boundary conditions. For critical infrastructure applications, such as structural health monitoring of bridges and buildings, particularly in digital twin contexts, low-latency energy-efficient topology optimization is essential. Traditionally, topology optimization relies on finite element analysis (FEA), a computationally intensive process. Recent advances in deep neural networks (DNNs) have introduced data driven alternatives to FEA, substantially reducing computation time while maintaining solution quality. These DNNs have complex architectures and implementing them on inference-class GPUs results in high latency and poor energy efficiency. To address this challenge, we present a hardware accelerated implementation of a topology optimization neural network (CRONet) on the AMD Versal AI Engine-ML (AIE-ML) architecture. Our approach efficiently exploits the parallelism and memory hierarchy of AIE-ML engines to optimize the execution of various neural network operators. We are the first to implement an end-to-end neural network fully realized on the AIE-ML array, where all intermediate activations and network weights reside on-chip throughout inference, eliminating any reliance on DRAM for intermediate data movement. Experimental results demonstrate that our implementation achieves up to 2.49x improvement in latency and up to 4.18x improvement in energy efficiency compared to an inference-class ML-optimized GPU in the same power budget (Nvidia T4) after scaling for technology node. These results highlight the potential of Versal AIE-ML based acceleration for enabling low-latency energy-efficient topology optimization.

• The paper presents the first complete mapping of CRONet’s hybrid CNN–RNN surrogate onto AMD Versal AIE-ML engines, keeping weights and intermediate activations on-chip while limiting DRAM traffic to inputs and outputs.
• The implementation combines three-level operator fusion, persistent local weights, scalable GEMM kernels, unrolled RNN dataflow, and congestion-aware placement to compile a 73%-utilized design in 8 minutes instead of failing after nearly an hour.
• The Versal VEK280 achieves 0.45–0.82 ms latency and 2.62–4.18× higher energy efficiency than a node-normalized NVIDIA T4 at batch size 1, although gains narrow for larger workloads and comparisons use BF16 on Versal versus FP32 on the T4.

Overview

This paper presents a hardware-accelerated implementation of CRONet, a convolutional-recurrent neural operator surrogate for topology optimization (TO), on the AMD Versal AI Engine-ML (AIE-ML) array. The authors claim, and to their knowledge correctly so, the first end-to-end mapping of a complete hybrid CNN–RNN network entirely onto the AIE-ML array: all weights and intermediate activations remain on-chip throughout inference, with DRAM traffic limited to input reads and output writes via GMIO. On the VEK280 board, the design achieves up to 2.49× lower latency and up to 4.18× higher energy efficiency than an NVIDIA T4 GPU at batch size 1, after normalizing T4 latency to the 7 nm process node.

Motivation and problem setting

Topology optimization determines material distribution within a design domain under load constraints, and its traditional workflow is bottlenecked by iterative finite element analysis (FEA). CRONet [(2604.14700)'s underlying model, from Olabiyi et al.] replaces FEA inside the optimization loop with a DeepONet-style dual-network surrogate: a BranchNet (a CRNN processing the evolving material layout) and a TrunkNet (a CNN processing the fixed load configuration), whose outputs combine into a predicted displacement field. Prior work reports up to 78% reduction in computational time on MBB-beam benchmarks with compliance within 5% of FEA solutions.

Two properties make CRONet difficult for conventional accelerators. First, its operator mix — 3D convolution, SiLU, 3D adaptive average pooling, RNN layers — is heterogeneous and irregularly shaped. Second, TO is inherently iterative: each iteration's output feeds the next, forcing batch size 1 and preventing GPUs from amortizing overhead across large batches. The authors argue this makes inference-class GPUs suboptimal, while FPGAs suffer low compute density per unit area; the AIE-ML dataflow architecture, with 304 VLIW vector engines (38×8 grid), 64 KB local memory per engine, and 512 KB Memory Tiles, offers fine-grained control suited to low-latency single-sample inference.

Architecture and methodology

The implementation maps both sub-networks as parallel dataflow pipelines on the AIE-ML array, deliberately avoiding the programmable logic entirely. This leaves the PL free for complementary edge tasks (sensor preprocessing, actuator driving) and keeps external bandwidth minimal enough that GMIO suffices without PLIO.

Three fusion strategies minimize data movement:

• L1 fusion: element-wise operations (e.g., SiLU) are fused directly into convolution and GEMM kernels.
• L2 fusion: outputs stream between adjacent subgraphs over direct AIE-to-AIE buffer connections, bypassing Memory Tiles and DRAM.
• L3 fusion: intermediate activations exceeding local memory capacity, or requiring reshaping, are staged through Memory Tiles rather than DRAM.

The GEMM kernel extends GAMA by persisting weights in AIE local memory and fully parameterizing the M, K, N dimensions. To escape GAMA's constraint that K fit within the 38-column cascade chain, large K dimensions are sliced across independent AIE clusters with partial sums reduced via an adder tree, decoupling scalability from the cascade limit. The RNN is implemented without a dedicated kernel: it is fully unrolled across time steps as a static dataflow graph, each step reusing GEMM kernels with Tanh fused into an add kernel via L1 fusion. Custom kernels are also provided for 2D/3D convolution (vectorized sliding-window along width), MaxPool2D, adaptive average pooling (2D/3D), and LUT-based SiLU. All operators are template-parameterized, forming a reusable, network-agnostic library intended for open-source release.

A notable systems contribution is congestion-aware placement. At 73% AIE utilization (223/304 engines), the default compiler placer fails to route the design even at maximum optimization level, after 48–62 minutes of compilation. The custom algorithm co-locates dataflow-adjacent kernels to shorten wires and relieve congestion, achieving successful compilation in 8 minutes. This is a strong practical result: without it, the full-scale design would not compile at all.

Quantization

Inference precision was evaluated against solution accuracy in a 100-iteration hybrid NN-FEA loop where CRONet is invoked only when prior-iteration error falls below threshold. BF16 matches FP32 exactly (100% accuracy, CRONet invoked for 33/100 iterations), whereas INT8 degrades accuracy to 90.91%. BF16 was therefore selected, aligning with native AIE-ML datatypes. This choice has a methodological consequence acknowledged by the authors: the T4 lacks efficient FP16/BF16 support at batch size 1, so the GPU baseline runs FP32 — reflecting realistic deployment constraints rather than peak GPU capability, but also meaning the comparison is not precision-matched.

Results

For the medium workload (30×20), Versal achieves 0.52 ms latency at 21.15 W versus 1.19 ms (node-normalized) at 35 W on the T4 — a 2.29× latency improvement and 3.79× energy-efficiency improvement. Across workloads:

Size	VEK280 latency / power	T4 latency* / power	Latency gain	Energy-eff. gain
30×10	0.45 ms / 19.06 W	1.12 ms / 32 W	2.49×	4.18×
30×20	0.52 ms / 21.15 W	1.19 ms / 35 W	2.29×	3.79×
60×20	0.82 ms / 21.49 W	1.25 ms / 37 W	1.52×	2.62×

*T4 latencies normalized to 7 nm per Stillmaker et al.

Versal latency scales sub-linearly (only 1.82× for a 4× input growth), which the authors attribute to parallel execution and on-chip fusion absorbing added computation. Power remains nearly flat (19–21.5 W). The advantage narrows at larger sizes because the T4 gains utilization as computation grows — an honest caveat on the headline numbers.

Layer-wise profiling identifies BranchNet's CONV2D as the dominant bottleneck at 55.3% of execution time, due to irregular memory access patterns limiting VLIW instruction-width utilization; AAP3D contributes 18.1%. These two layers are explicitly left as future optimization targets. Resource allocation shows BranchNet consuming 134 engines versus TrunkNet's 78; despite the imbalance, both run concurrently, with TrunkNet finishing early.

Limitations and open questions

The paper concedes several constraints plainly. Architecturally, GMIO bandwidth and port count are limited, and Memory Tile buffer descriptors restrict the complexity of data transformations that can be staged between layers. Programmability is a more fundamental limitation: the ADF graph model enforces rigid, synchronous dataflow with fixed iteration counts per graph, precluding selective per-subgraph iteration control and runtime reconfiguration; statically allocated resources cannot be reused across graphs, causing underutilization. The comparison itself carries caveats: the GPU baseline is not precision-matched (FP32 vs. BF16), power is measured differently on each platform (PMIC rails vs. nvidia-smi), and the advantage shrinks monotonically with problem size, leaving unclear at what workload scale the T4 would close or reverse the gap. Finally, the reusable operator library currently covers only CRONet's operators, and improved memory-access patterns for CONV2D and AAP3D remain unexplored.

Conclusion

This paper demonstrates a complete, DRAM-free-on-intermediates deployment of a hybrid CNN–RNN neural operator on the AMD Versal AIE-ML array, enabled by three-level fusion, persistent on-chip weights, extended GEMM mappings that decouple scaling from the cascade chain limit, and a congestion-aware placement strategy without which the design does not compile. Measured results show consistent latency and energy-efficiency advantages over an inference-class GPU at batch size 1, with the margin narrowing as workloads grow. The main open questions are whether the placement and fusion techniques generalize beyond CRONet's operator set, and whether programming-model extensions for per-subgraph iteration control could recover the utilization currently lost to static resource allocation.