[Paper Review] Reinforcement Learning Approach for Mapping Applications to Dataflow-Based Coarse-Grained Reconfigurable Array

Andre Xian Ming Chang, Parth Khopkar+7 more|arXiv (Cornell University)|May 26, 2022

Parallel Computing and Optimization Techniques4 citations

TL;DR

This paper proposes a reinforcement learning (RL) framework with Global Graph Attention (GGA) and output masking to automatically map dataflow applications to a Coarse-Grained Reconfigurable Array (CGRA) Streaming Engine. The method uses Proximal Policy Optimization (PPO) to learn optimal instruction scheduling, achieving 10% better clock cycle performance and 20% higher reward with masking, significantly outperforming baseline methods in sparse search spaces.

ABSTRACT

The Streaming Engine (SE) is a Coarse-Grained Reconfigurable Array which provides programming flexibility and high-performance with energy efficiency. An application program to be executed on the SE is represented as a combination of Synchronous Data Flow (SDF) graphs, where every instruction is represented as a node. Each node needs to be mapped to the right slot and array in the SE to ensure the correct execution of the program. This creates an optimization problem with a vast and sparse search space for which finding a mapping manually is impractical because it requires expertise and knowledge of the SE micro-architecture. In this work we propose a Reinforcement Learning framework with Global Graph Attention (GGA) module and output masking of invalid placements to find and optimize instruction schedules. We use Proximal Policy Optimization in order to train a model which places operations into the SE tiles based on a reward function that models the SE device and its constraints. The GGA module consists of a graph neural network and an attention module. The graph neural network creates embeddings of the SDFs and the attention block is used to model sequential operation placement. We show results on how certain workloads are mapped to the SE and the factors affecting mapping quality. We find that the addition of GGA, on average, finds 10% better instruction schedules in terms of total clock cycles taken and masking improves reward obtained by 20%.

Motivation & Objective

• Address the NP-hard, sparse, and complex problem of manually mapping Synchronous Data Flow (SDF) graphs to Streaming Engine (SE) tiles with architectural constraints.
• Reduce the need for expert knowledge and manual effort in SE application mapping by enabling automated, unsupervised learning-based optimization.
• Improve mapping quality and sample efficiency in a sparse search space by integrating attention mechanisms and action masking.
• Enable generalization across diverse workloads through pre-training and fine-tuning on varied IR graphs.
• Integrate the RL mapper into the SE toolchain to accelerate development and support automated exploration of trade-offs in scheduling.

Proposed method

• Employ Proximal Policy Optimization (PPO) as the RL algorithm to train a policy network that selects tile and slot placements for SDF nodes.
• Introduce a Global Graph Attention (GGA) module combining a graph neural network (GNN) for SDF embedding and a multi-head attention mechanism to model node dependencies during placement.
• Apply output masking to eliminate invalid actions (e.g., over-constrained or data-dependent placements) at each step, reducing search space and improving training stability.
• Use a simulated SE environment as the RL environment, modeling tile memory, dataflow dependencies, and inter-tile latency to compute rewards.
• Implement sorted iteration order (topological) to improve training by placing dependent nodes only after their predecessors.
• Apply pre-training on random IR graphs followed by fine-tuning on specific workloads like FFT to transfer knowledge and improve convergence.

Experimental results

Research questions

• RQ1 Can a deep reinforcement learning framework effectively explore the vast and sparse search space of instruction mappings on a CGRA without manual supervision?
• RQ2 How does incorporating a graph attention mechanism (GGA) improve the RL agent’s ability to reason about dataflow dependencies and placement quality?
• RQ3 To what extent does output masking of invalid actions enhance learning efficiency and final reward in the SE mapping task?
• RQ4 How does the order of node iteration (e.g., topological vs. random) affect the success rate and quality of the learned mappings?
• RQ5 Can pre-trained RL models be effectively fine-tuned on new workloads to achieve faster convergence and better performance than training from scratch?

Key findings

• The addition of the Global Graph Attention (GGA) module improved the best schedule found by 10% in terms of total clock cycles across various graph complexities and SE configurations.
• Output masking increased the reward obtained by 20% compared to using negative rewards for invalid actions, significantly improving sample efficiency and convergence.
• Using topological iteration order for node processing led to higher rewards and more stable training, as it aligned with dataflow semantics and reduced prediction burden.
• Pre-training on random IR graphs followed by fine-tuning on specific workloads like FFT resulted in higher initial and final rewards than training from scratch, demonstrating transfer learning potential.
• The RL mapper successfully found valid and optimized mappings for diverse workloads, including FFT and distance calculation, even on a 64-tile SE configuration, outperforming simulated annealing in sparse search spaces.
• The framework demonstrated scalability and adaptability, showing promise for integration into the SE toolchain to reduce manual effort and accelerate application deployment.