[challenge]: Contraction Path Search Incorporating Permutation Overhead

Released by

Shi-Xin Zhang, IOP-CAS

Contact email

shixinzhang@iphy.ac.cn

Method

Quantum Circuit Simulation

Challenge issue

Background

Tensor network contraction is a fundamental computational kernel for quantum simulation and many-body physics. Currently, modern tensor network frameworks separate the problem into two distinct layers:

Contraction Planning: Optimizers (like cotengra or opt_einsum) determine the contraction order, prioritizing minimum FLOPs and peak memory.

Execution: Backends (like JAX, cuTensorNet, or cuBLAS) translate these contractions into GPU operations, handling tensor layouts and memory movement.

This hardware-agnostic separation has been highly successful. However, with the advent of modern AI-oriented GPUs featuring immensely powerful Tensor Cores, this abstraction layer has some new aspect to consider.

The Problem: The Hidden Cost of Layouts

Traditional contraction planners rely on an algebraic cost model: Cost ≈ FLOPs + alpha Memory.

On modern GPUs, this model is incomplete. A standard tensor contraction typically requires: Permutation -> GEMM -> Permutation. While the GEMM operation is aggressively accelerated by Tensor Cores, permutations and layout conversions are memory-bound, limited by HBM bandwidth, and cannot utilize Tensor Cores.

Because the contraction planner ignores layouts, it might select a path with slightly fewer FLOPs (Strategy A) that requires massive, expensive permutations. Meanwhile, it rejects a path with slightly more FLOPs (Strategy B) that maintains favorable tensor layouts and avoids costly memory movement. As compute becomes cheaper and data movement remains expensive, two contraction paths with identical FLOPs can have drastically different real-world runtimes.

The Challenge

Can we extend tensor network contraction planning from simply finding the best contraction order to jointly optimizing the contraction order and tensor layout?

The challenge is to design a tree search scheme that explicitly considers permutation overheads during the contraction planning phase. Instead of treating permutations as zero-cost operations, the optimizer should formulate a schedule comprising:

The contraction order.

Intermediate tensor layouts.

The necessary data movement strategy.

Potential Directions & Objectives

Participants are free to tackle this problem creatively, but a successful solution will likely involve:

A New Cost Model: Formulate a cost function that penalizes expensive data movement alongside standard GEMM FLOPs and peak memory (e.g., Cost=α⋅FLOPs+β⋅Data Movement+…).

Layout-Aware Search: Modify the search algorithm (e.g., extending cotengra) so intermediate nodes track not just shapes and indices, but layout and stride information.

Benchmarking & Profiling: Demonstrate your optimizer's effectiveness. Show that by intentionally selecting paths with fewer permutations (even if FLOPs are slightly higher), your scheme achieves measurable end-to-end speedups on modern hardware compared to naive, FLOP-centric optimizers.

Outcome: A successful implementation could redefine how next-generation tensor network compilers approach optimization on modern GPU architectures.

No files, tests, or entry points are named. Start by locating the contraction-planning implementation and its existing cost model, using cotengra as a possible reference. Done means producing a search that accounts for contraction order, intermediate layouts, and data movement, then benchmarking it against FLOP-centric planning on modern hardware.