IP Library › Granted Patent US 12,461,723
Granted Patent B2
US 12,461,723 · App. 17/921,933 · Granted Nov 4, 2025

Learned graph optimizations for compilers

Inventors: Yanqi Zhou (Sunnyvale, CA); Sudip Roy (Sunnyvale, CA); Amirali Abdolrashidi (Riverside, CA); Daniel Lin-Kit Wong (Pittsburgh, PA); Chao Ma (Mountain View, CA); Qiumin Xu (Santa Clara, CA); Hanxiao Liu (Santa Clara, CA); Phitchaya Mangpo Phothilimthana (Mountain View, CA); Shen Wang (Sunnyvale, CA); Anna Darling Goldie (San Francisco, CA); Azalia Mirhoseini (Mountain View, CA); James Laudon (Madison, WI)
Assignee: Google LLC
G06F8/443G06F8/451G06N3/044
View Patent ↗
Loading inventors, assignments & file history…
Monitor This Case
Get email alerts when status or documents change.
Order Certified Copies
Most orders are placed with the USPTO same day — all within 24 business hours.
Order via The Patent Place →
Pre-filled with this patent's details
Quick Facts
Patent No.
US 12,461,723
App. No.
17/921,933
Granted
Nov 4, 2025
Kind
B2
Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for compiler optimizations using a compiler optimization network. One of the methods includes receiving an input program, wherein the input program defines a graph of operation modules, wherein each node in the graph is a respective operation module, and each edge between nodes in the graph represents one operation module receiving the output generated by another operation module. The input program is processed by a compiler optimization network comprising a graph-embedding network that is configured to encode operation features and operation dependencies of the operation modules of the input program into a graph embedding representation and a policy network that is configured to generate an optimization action for each of one or more nodes encoded in the graph embedding representation. The compiler optimization network generates an output optimization plan comprising one or more optimization actions for the input program.

Claims (42)

1 . A computer-implemented method performed by a system comprising one or more computers and configured to execute a compiler optimization network, the method comprising:

receiving an input program, wherein the input program defines a graph of operation modules, wherein each node in the graph is a respective operation module, and each edge between nodes in the graph represents one operation module receiving the output generated by another operation module;

providing a representation of the input program to the compiler optimization network comprising:

a graph-embedding network executable by the system that is configured to encode operation features and operation dependencies of the operation modules of the input program into a graph embedding representation, and

a policy network executable by the system that is configured to generate an optimization action for each of one or more nodes encoded in the graph embedding representation, wherein the policy network employs segmented recurrent attention layers to concurrently generate optimization actions for multiple different tasks, wherein each layer of the segmented recurrent attention layers is dedicated to generating a respective optimization action for a different task of the multiple different tasks;

obtaining, from the compiler optimization network, an output optimization plan comprising one or more optimization actions for each of the multiple different tasks of the input program; and

executing the input program using the output optimization plan.

2 . The method of claim 1 , wherein executing the input program using the output optimization plan further comprises:

scheduling the input program for execution on the one or more devices using the output optimization plan generated by the compiler optimization network.

3 . The method of claim 1 , further comprising performing a graph-embedding process using the graph-embedding network comprising iteratively aggregating respective feature representations for one or more neighboring nodes of each node in the graph.

4 . The method of claim 1 , wherein the policy network includes a feature modulation layer that is configured to re-weight network parameters.

5 . The method of claim 1 , further comprising jointly training the compiler optimization network over a set of N different dataflow graphs.

6 . The method of claim 5 , wherein the N different dataflow graphs represent programs for training machine learning models in different machine-learning domains.

7 . The method of claim 1 , wherein training the compiler optimization network comprises training the compiler optimization network end-to-end over both the graph-embedding network and the policy network.

8 . The method of claim 1 , wherein the policy network generates recurrent actions for multiple using residual connections and parameter sharing across multiple recurrent attention layers.

9 . The method of claim 1 , wherein the compiler optimization network implements a policy that specifies whether two operation modules should be fused.

10 . The method of claim 1 , wherein the compiler optimization network implements a policy that assigns placement of operation modules to particular devices.

11 . The method of claim 1 , wherein the compiler optimization network implements a policy that assigns a scheduling priority for a plurality of operation modules in the graph.

12 . A system configured to execute a compiler optimization network, the system comprising:

one or more computers and one or more storage devices storing instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform operations comprising:

receiving an input program, wherein the input program defines a graph of operation modules, wherein each node in the graph is a respective operation module, and each edge between nodes in the graph represents one operation module receiving the output generated by another operation module;

providing a representation of the input program to the compiler optimization network comprising:

a graph-embedding network executable by the system that is configured to encode operation features and operation dependencies of the operation modules of the input program into a graph embedding representation, and

a policy network executable by the system that is configured to generate an optimization action for each of one or more nodes encoded in the graph embedding representation, wherein the policy network is configured to employ segmented recurrent attention layers to concurrently generate optimization actions for multiple different tasks, wherein each layer of the segmented recurrent attention layers is dedicated to generating a respective optimization action for a different task of the multiple different tasks;

obtaining, from the compiler optimization network, an output optimization plan comprising one or more optimization actions for each of the multiple different tasks of the input program; and

executing the input program using the output optimization plan.

13 . The system of claim 12 , wherein executing the input program using the output optimization plan further comprises:

scheduling the input program for execution on the one or more devices using the output optimization plan generated by the compiler optimization network.

14 . The system of claim 12 , wherein the operations further comprise performing a graph-embedding process using the graph-embedding network comprising iteratively aggregating respective feature representations for one or more neighboring nodes of each node in the graph.

15 . The system of claim 12 , wherein the policy network includes a feature modulation layer that is configured to re-weight network parameters.

16 . The system of claim 12 , wherein the operations further comprise jointly training the compiler optimization network over a set of N different dataflow graphs.

17 . The system of claim 16 , wherein the N different dataflow graphs represent programs for training machine learning models in different machine-learning domains.

18 . One or more non-transitory computer storage media encoded with computer program instructions that when executed by one or more computers cause the one or more computers to perform operations for a system comprising one or more computers and configured to execute a compiler optimization network, the operations comprising:

receiving an input program, wherein the input program defines a graph of operation modules, wherein each node in the graph is a respective operation module, and each edge between nodes in the graph represents one operation module receiving the output generated by another operation module;

providing a representation of the input program to the compiler optimization network comprising:

a graph-embedding network executable by the system that is configured to encode operation features and operation dependencies of the operation modules of the input program into a graph embedding representation, and

a policy network executable by the system that is configured to generate an optimization action for each of one or more nodes encoded in the graph embedding representation, wherein the policy network is configured to employ segmented recurrent attention layers to concurrently generate optimization actions for multiple different tasks, wherein each layer of the segmented recurrent attention layers is dedicated to generating a respective optimization action for a different task of the multiple different tasks;

obtaining, from the compiler optimization network, an output optimization plan comprising one or more optimization actions for each of the multiple different tasks of the input program; and

executing the input program using the output optimization plan.

19 . The one or more non-transitory computer storage media of claim 18 , wherein executing the input program using the output optimization plan further comprises:

scheduling the input program for execution on the one or more devices using the output optimization plan generated by the compiler optimization network.

20 . The one or more non-transitory computer storage media of claim 18 , the operations further comprising performing a graph-embedding process using the graph-embedding network comprising iteratively aggregating respective feature representations for one or more neighboring nodes of each node in the graph.

Assignments (1)
ASSIGNMENT OF ASSIGNOR'S INTEREST Recorded Apr 25, 2023
From: ZHOU, YANQI; ROY, SUDIP; ABDOLRASHIDI, AMIRALI; WONG, DANIEL LIN-KIT; MA, CHAO; XU, QIUMIN; LIU, HANXIAO; PHOTHILIMTHANA, PHITCHAYA MANGPO; WANG, SHEN; GOLDIE, ANNA DARLING; MIRHOSEINI, AZALIA; LAUDON, JAMES
To: GOOGLE LLC
Reel/Frame 063427/0020 →
Continuity (2)
Provisional Application 63035640 · Jun 5, 2020
Related Publication 20230176840A1 · Jun 8, 2023
References Cited (62)
US 11169786B2 · Badlani · 2021 [cited by examiner]
US 12093806B1 · Zejda · 2024 [cited by examiner]
US 20190391796A1 · Brady et al. · 2019 [cited by applicant]
US 20200089755A1 · Shazeer · 2020 [cited by examiner]
US 20200104681A1 · Li · 2020 [cited by examiner]
US 20200249998A1 · Che · 2020 [cited by examiner]
US 20200293838A1 · Li · 2020 [cited by examiner]
US 20210133591A1 · Elango · 2021 [cited by examiner]
US 20210182036A1 · Shafiq · 2021 [cited by examiner]
US 20210192314A1 · Aarts · 2021 [cited by examiner]
US 20210232873A1 · Kothari · 2021 [cited by examiner]
US 20210286831A1 · Girardi · 2021 [cited by examiner]
CN 111126386A · 2020 [cited by examiner]
WO WO2019018332A1 · 2019 [cited by examiner]
Aditya Paliwal, Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs, Feb. 2020, pp. 1-24. https://arxiv.org/pdf/1905.02494 (Year: 2020). [cited by examiner]
Zhilin Yang, Transformer-XL: Unleashing the Potential of Attention Models, 2019, pp. 1-6. https://research.google/blog/transformer-xl-unleashing-the-potential-of-attention-models/ (Year: 2019). [cited by examiner]
English translation, Goldie (WO 2019018332 A1), 2019, pp. 1-14. (Year: 2019). [cited by examiner]
Yao Wan, Improving Automatic Source Code Summarization via Deep Reinforcement Learning, 2018, pp. 1-12. https://arxiv.org/pdf/1811.07234 (Year: 2018). [cited by examiner]
Chris Cummins, PROGRAML: Graph-Based Deep Learning for Program Optimization and Analysis, pp. 1-20, 2020. chrome-extension://efaidnbmnnnibpcajpcglclefindmkaj/https://arxiv.org/pdf/2003.10536 (Year: 2020). [cited by examiner]
English translation, Zhou et al. (CN 111126386 A), 2020, pp. 1-9. (Year: 2020). [cited by examiner]
Adams et al., “Learning to optimize halide with tree search and random programs,” ACM Transactions on Graphics, Jul. 12, 2019, 38(4):1-12. [cited by applicant]
Addanki et al. “Placeto: Learning generalizable device placement algorithms for distributed machine learning,” CoRR, Submitted on Jun. 20, 2019, ar Xiv: 1906.08879v1, 15 pages. [cited by applicant]
Ba et al., “Layer normalization,” CoRR, Submitted on Jul. 21, 2016, arXiv: 1607.06450v1, 14 pages. [cited by applicant]
Chen et al., “Learning to optimize tensor programs,” CoRR, Submitted on May 21, 2018, arXiv: 1805.08166v1, 17 pages. [cited by applicant]
Chen et al., “Learning to perform local rewriting for combinatorial optimization,” CoRR, Submitted on Oct. 30, 2019, arXiv: 1810.00337v5, 20 pages. [cited by applicant]
Cheung et al., “Superposition of many models into one,” CoRR, Submitted on Jun. 17, 2019, arXiv:1902.05522v2, 18 pages. [cited by applicant]
Ciosek et al., “Knossos: Compiling AI with AI,” Proceedings of International Conference on Learning Representations, Sep. 25, 2019, 23 pages. [cited by applicant]
Dai et al., “Transformer-xl: Attentive language models beyond a fixed-length context,” Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, Jul. 28-Aug. 2, 2019, pp. 2978-2988. [cited by applicant]
Dai, “Improving deep generative modeling with applications,” Thesis for the degree of Doctor of Philosophy, Carnegie Mellon University, School of Computer Science, Sep. 19, 2019, 131 pages. [cited by applicant]
Devlin et al., “BERT: Pre-training of deep bidirectional transformers for language understanding,” CoRR, Submitted on Oct. 11, 2018, arXiv:1810.04805v1, 14 pages. [cited by applicant]
Gao et al., “Spotlight: Optimizing device placement for training deep neural networks,” Proceedings of the 35th International Conference on Machine Learning (PMLR), Jul. 10-15, 2018, 80:1676-1684. [cited by applicant]
Google for Developers [online], “XLA—TensorFlow, compiled,” Mar. 6, 2017, retrieved on Feb. 16, 2024, retrieved from URL<https://developers.googleblog.com/2017/03/xla-tensorflow-compiled.html>, 7 pages. [cited by applicant]
Hamilton et al., “Inductive representation learning on large graphs,” Proceedings of the 31st International Conference on Neural Information Processing Systems, Dec. 2017, pp. 1025-1035. [cited by applicant]
Hestness et al., “Deep learning scaling is predictable, empirically,” CoRR, Submitted on Dec. 1, 2017, arXiv:1712.00409v1, 19 pages. [cited by applicant]
Huang et al., “Gpipe: Efficient training of giant neural networks using pipeline parallelism,” CoRR, Submitted on Nov. 20, 2018, arXiv:1811.06965v3, 10 pages. [cited by applicant]
International Preliminary Report on Patentability in International Appln. No. PCT/US2021/036250, mailed on Dec. 15, 2022, 8 pages. [cited by applicant]
International Search Report and Written Opinion in International Appln. No. PCT/US2021/036250, mailed on Sep. 22, 2021, 15 pages. [cited by applicant]
Jia et al., “Beyond data and model parallelism for deep neural networks,” Proceedings of Machine Learning and Systems, Apr. 15, 2019, 1:1-13. [cited by applicant]
Jia et al., “TASO: optimizing deep learning computation with automatic generation of graph substitutions,” Proceedings of the 27th ACM Symposium on Operating Systems Principles, Oct. 27, 2019, pp. 47-62. [cited by applicant]
Jozefowicz et al., “Exploring the limits of language modeling,” CoRR, Submitted on Feb. 11, 2016, arXiv:1602.02410v2, 11 pages. [cited by applicant]
Karypis et al., “A fast and high quality multilevel scheme for partitioning irregular graphs,” SIAM Journal on scientific Computing, Dec. 1998, 20(1):359-392. [cited by applicant]
Lattner et al., “MLIR Primer: A Compiler Infrastructure for the End of Moore's Law,” Proceeding of International Symposium on Code Generation and Optimization (CGO), Feb. 16-20, 2019, 45 pages. [cited by applicant]
Mahajan et al., “Exploring the limits of weakly supervised pretraining,” Proceedings of the European Conference on Computer Vision (ECCV), Sep. 8-14, 2018, pp. 185-201. [cited by applicant]
Mirhoseini et al., “A hierarchical model for device placement,” Proceeding of International Conference on Learning Representations, Feb. 15, 2018, 11 pages. [cited by applicant]
Mirhoseini et al., “Chip placement with deep reinforcement learning,” CoRR, Submitted on Apr. 22, 2020, arXiv:2004.10746v1, 15 pages. [cited by applicant]
Mirhoseini et al., “Device placement optimization with reinforcement learning,” Proceeding of International Conference on Machine Learning (PMLR), Jul. 17, 2017, pp. 2430-2439. [cited by applicant]
Narayanan et al., “PipeDream: Generalized pipeline parallelism for DNN training,” Proceedings of the 27th ACM Symposium on Operating Systems Principles, Oct. 2019, pp. 1-15. [cited by applicant]
Office Action in European Appln No. 21743313.5, mailed on Oct. 20, 2023, 4 pages. [cited by applicant]
Paliwal et al., “REGAL: Transfer learning for fast optimization of computation graphs,” CoRR, Submitted on May 2019, arXiv:1905.02494v2, 12 pages. [cited by applicant]
Radford et al., “Language models are unsupervised multitask learners,” OpenAI, Feb. 24, 2019, 24 pages. [cited by applicant]
Real et al., “Regularized evolution for image classifier architecture search,” CoRR, Submitted on Feb. 5, 2018, arXiv:1802.01548v1, 15 pages. [cited by applicant]
Rotem et al., “Glow: Graph lowering compiler techniques for neural networks,” CoRR, Submitted on May 4, 2018, arXiv: 1805.00907v2, 10 pages. [cited by applicant]
Schulman et al., “Proximal policy optimization algorithms,” CoRR, Submitted on Aug. 28, 2017, arXiv:1707.06347v2, 12 pages. [cited by applicant]
Shazeer et al., “Outrageously large neural networks: The sparsely-gated mixture-of-experts layer,” CoRR, Submitted on Jan. 23, 2017, arXiv:1701.06538v1, 19 pages. [cited by applicant]
Sutskever et al., “Sequence to sequence learning with neural networks,” CoRR, Submitted on Dec. 14, 2014, arXiv:1409.3215v3, 9 pages. [cited by applicant]
Szegedy et al., “Rethinking the inception architecture for computer vision,” CoRR, Submitted on Dec. 11, 2015, arXiv:1512.00567v3, 10 pages. [cited by applicant]
Van den Oord et al., “WaveNet: A generative model for raw audio,” CoRR, Sep. 19, 2016, arXiv:1609.03499v2, 15 pages. [cited by applicant]
Vaswani et al., “Attention is all you need,” Proceedings of the 31st International Conference on Neural Information Processing Systems, Dec. 2017, 11 pages. [cited by applicant]
Xu et al., “How powerful are graph neural networks?” Proceeding of International Conference on Learning Representations, May 6-9, 2019, 17 pages. [cited by applicant]
You et al., “Graph convolutional policy network for goal-directed molecular graph generation,” CoRR, Submitted on Nov. 18, 2018, arXiv:1806.02473v2, 12 pages. [cited by applicant]
Zaremba et al., “Recurrent neural network regularization,” CoRR, Submitted on Dec. 21, 2014, arXiv:1409.2329v4, 8 pages. [cited by applicant]
Zhou et al., “GDP: Generalized device placement for dataflow graphs,” CoRR, Submitted on Sep. 28, 2019, arXiv:1910.01578v1, 11 pages. [cited by applicant]