Skip to main content
Warning: You are using the test version of PyPI. This is a pre-production deployment of Warehouse. Changes made here affect the production instance of TestPyPI (
Help us improve Python packaging - Donate today!

Explore the energy-efficient dataflow scheduling for neural networks.

Project Description

Neural Network Dataflow Scheduling

This Python tool allows you to explore the energy-efficient dataflow scheduling for neural networks (NNs), including array mapping, loop blocking and reordering, and parallel partitioning.

For hardware, we assume an Eyeriss-style NN accelerator [Chen16], i.e., a 2D array of processing elements (PEs) with a local register file in each PE, and a global SRAM buffer shared by all PEs. We further support a tiled architecture with multiple nodes that can partition and process the NN computations in parallel. Each node is an Eyeriss-style engine as above.

In software, we decouple the dataflow scheduling into three subproblems:

  • Array mapping, which deals with mapping one 2D convolution computation (one 2D ifmap convolves with one 2D filter to get one 2D ofmap) onto the hardware PE array. We support row stationary mapping [Chen16].
  • Loop blocking and reordering, which decides the order between all 2D convolutions by blocking and reordering the nested loops. We support exhaustive search over all blocking and reordering schemes [Yang16], and analytical bypass solvers [Gao17].
  • Partitioning, which partitions the NN computations for parallel processing. We support batch partitioning, fmap partitioning, output partitioning, input partitioning, and the combination between them (hybrid) [Gao17].

See the details in our ASPLOS‘17 paper [Gao17].

If you use this tool in your work, we kindly request that you reference our paper(s) below, and send us a citation of your work.

  • Gao et al., “TETRIS: Scalable and Efficient Neural Network Acceleration with 3D Memory”, in ASPLOS, April 2017 [Gao17].


First, define the NN structure in nn_dataflow/nns. We already defined several popular NNs for you, including AlexNet, VGG-16, GoogLeNet, ResNet-152, etc.

Then, use nn_dataflow/tools/ to search for the optimal dataflow for the NN. For detailed options, type:

> python ./nn_dataflow/tools/ -h

You can specify NN batch size and word size, PE array dimensions, number of tile nodes, register file and global buffer capacity, and the energy cost of all components. Note that, the energy cost of array bus should be the average energy of transferring the data from the buffer to one PE, not local neighbor transfer; the unit static energy cost should be the static energy of one node in one clock cycle.

Other options include:

  • --mem-type: 2D or 3D. With 2D memory, memory channels are only on the left and right sides of the chip; with 3D memory, memory channels are on the top of all tile nodes (one per each).
  • --disable-bypass: a combination of i, o, f, whether to disallow global buffer bypass for ifmaps, ofmaps, and weights.
  • --solve-loopblocking: whether to use analytical bypass solvers for loop blocking and reordering. See [Gao17].
  • --hybrid-partitioning: whether to use hybrid partitioning in [Gao17]. If not enabled, use naive partitioning, i.e., fmap partitioning for CONV layers, and output partitioning for FC layers.
  • --batch-partitioning and --ifmap-partitioning: whether the hybrid partitioning also explores batch and input partitioning.

Code Structure

  • nn_dataflow
    • core
      • Top-level dataflow exploration: nn_dataflow, nn_dataflow_scheme.
      • Layer scheduling: scheduling.
      • Array mapping: map_strategy.
      • Loop blocking and reordering: loop_blocking, loop_blocking_scheme, loop_blocking_solver.
      • Partitioning: partition, partition_scheme.
      • Network and layer: network, layer.
    • nns: example NN definitions.
    • tests: unit tests.
    • tools: executables.

Verification and Testing

To verify the tool against the Eyeriss result [Chen16], see nn_dataflow/tests/dataflow_test/

To run (unit) tests, do one of the following:

> python -m unittest discover

> python -m pytest

> pytest

To check code coverage with pytest-cov plug-in:

> pytest --cov=nn_dataflow


[Gao17](1, 2, 3, 4, 5, 6) Gao, Pu, Yang, Horowitz, and Kozyrakis, TETRIS: Scalable and Efficient Neural Network Acceleration with 3D Memory, in ASPLOS. April, 2017.
[Chen16](1, 2, 3) Chen, Emer, and Sze, Eyeriss: A Spatial Architecture for Energy-Efficient Dataflow for Convolutional Neural Networks, in ISCA. June, 2016.
[Yang16]Yang, Pu, Rister, Bhagdikar, Richardson, Kvatinsky, Ragan-Kelley, Pedram, and Horowitz, A Systematic Approach to Blocking Convolutional Neural Networks, arXiv preprint, 2016.
Release History

Release History

This version
History Node


Download Files

Download Files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

File Name & Checksum SHA256 Checksum Help Version File Type Upload Date
nn_dataflow-1.5-py2-none-any.whl (153.5 kB) Copy SHA256 Checksum SHA256 py2 Wheel Aug 23, 2017
nn_dataflow-1.5.tar.gz (86.2 kB) Copy SHA256 Checksum SHA256 Source Aug 23, 2017

Supported By

WebFaction WebFaction Technical Writing Elastic Elastic Search Pingdom Pingdom Monitoring Dyn Dyn DNS Sentry Sentry Error Logging CloudAMQP CloudAMQP RabbitMQ Heroku Heroku PaaS Kabu Creative Kabu Creative UX & Design Fastly Fastly CDN DigiCert DigiCert EV Certificate Rackspace Rackspace Cloud Servers DreamHost DreamHost Log Hosting