# LanPose: Language-Instructed 6D Object Pose Estimation
for Robotic Assembly

Bowen Fu1, Sek Kun Leong1, Yan Di2, Jiwen Tang1 and Xiangyang Ji1 1Department
of Automation and BNRist, Tsinghua University, Beijing, China. 2Technical
University of Munich (TUM), Munich, Germany.

###### Abstract

Comprehending natural language instructions is a critical skill for robots to
cooperate effectively with humans. In this paper, we aim to learn 6D poses for
robotic assembly by natural language instructions. For this purpose, Language-
Instructed 6D Pose Regression Network (LanPose) is proposed to jointly predict
the 6D poses of the observed object and the corresponding assembly position.
Our proposed approach is based on the fusion of geometric and linguistic
features, which allows us to finely integrate multi-modality input and map it
to the 6D pose in $SE(3)$ space by the cross-attention mechanism and the
language-integrated 6D pose mapping module, respectively. To validate the
effectiveness of our approach, an integrated robotic system is established to
precisely and robustly perceive, grasp, manipulate and assemble blocks by
language commands. 98.09 and 93.55 in ADD(-S)-0.1d are derived for the
prediction of 6D object pose and 6D assembly pose, respectively. Both
quantitative and qualitative results demonstrate the effectiveness of our
proposed language-instructed 6D pose estimation methodology and its potential
to enable robots to better understand and execute natural language
instructions.

##  I Introduction

Assembling blocks by language instructions is quite spontaneous for kids. The
target blocks are first identified and manipulated, and then assembled with
tight tolerances. In this process, the natural language instructions, e.g.,
insert to and stack on, need to be comprehended and mapped to a specific
action. One approach to enable robots to conduct similar tasks is through the
integration of geometric and linguistic information, which allows for the
prediction of a precise presentation in 3D space, specifically, the 6D object
pose. In this work, we focus on learning 6D object poses from natural language
instructions to assemble blocks by a robotic system.

Thanks to the integrative development of computer vision (CV) and natural
language processing (NLP), the fusion of image and natural language
information is possible. Specifically, visual grounding tasks [1, 2, 3]
provide the 2D localization information referring to the language expression.
Several works introduce visual grounding to robotic manipulation tasks. [4, 5]
both perform human-robot interaction by natural language instructions. On the
other hand, [6] emphasizes the history-dependent manipulation to preferably
conduct pick-and-place tasks. However, the objects are merely located by 2D
center points or 2D bounding boxes. Without the information in 3D space with
full 6 degrees of freedom, the precise assembly operation cannot be executed.

The key challenge lies in reasoning the 6D object pose from both image and
natural language information. The geometric features reflect the localization
and orientation of the perceived target object, which contains the surrounding
context information of the assembly position. Meanwhile, the linguistic
features reveal the desired interaction between two objects, in which the
relative relationship can be obtained. Recently, [7] proposes to learn to
grasp category-level objects by a sentence that describes the location of the
objects. Though 6D poses are predicted for grasping, the image-text
integration is only at 2D detection stage. Thus, only 2D location information
is extracted from the text, rendering the higher-level interaction between
objects impossible. To accomplish high-precision robotic assembly tasks, the
network has to extract and integrate the image and text features, and map them
to the 6D pose in $SE(3)$ space.

Figure 1: Example of our framework. Given a natural language instruction, (a)
the 2D localization of target objects and (b) the 6D poses of the target
objects and corresponding assembly position are predicted. (c) Our integrated
robotic assembly system.

In this work, we propose Language-Instructed 6D Pose Regression Network
(LanPose) to predict the 6D assembly position pose and establish an integrated
robotic system to assemble blocks by a single-sentence natural language
instruction (Fig. 1). Fed with the target object localization from 2D visual
grounding network, the LanPose integrates geometric and linguistic features by
an attention mechanism to simultaneously estimate the 6D poses of the observed
object and the target assembly position. Subsequently, an integrated robotic
system is established to precisely and robustly grasped, manipulated and
assembled the target object to the instructed position. Our networks take only
RGB input and are solely trained with synthetic data, which avoids time-
consuming real-world annotations, rendering our system extensible and
flexible. We achieve 98.09 and 93.55 in ADD(-S)-0.1d in 6D pose estimation for
perceived objects and assembly position and a 82.1% success rate in real-world
robotic assembly experiments.

Our contributions can be summarized as follows:

  * •

We present Language-Instructed 6D Pose Regression Network (LanPose) to jointly
predict high-precision 6D poses of the observed object and the corresponding
assembly position, which derives 98.09 and 93.55 in ADD(-S)-0.1d,
respectively.

  * •

We propose the intermediate geometric features fusion and the language-
integrated 6D pose mapping module, which integrates the geometric and
linguistic information by attention mechanism and precisely maps them to the
6D pose in $SE(3)$ space.

  * •

We further establish an integrated robotic system to precisely and robustly
assemble blocks by natural language instructions with a 82.1% success rate.

##  II Related Work

###  II-A 6D Object Pose Estimation

6D object pose estimation has received much attention in both computer vision
and robotics communities. A branch of approaches relies on establishing 2D-3D
correspondences followed by RANSAC-based P$n$P algorithm to indirectly predict
the 6D pose [8, 9, 10, 11] while another branch directly regresses the 6D pose
[12, 13, 14, 15, 16, 17]. In robotic manipulation tasks, [18, 19, 20, 21]
leverage 6D object pose information to grasp objects. [22] further utilize the
interaction between the robot and objects to self-supervised collecting data
for training 6D pose estimation network. Similarly, [23] proposes an iterative
self-training framework for robotic bin-picking. [24] learns task-relevant
grasping for industrial objects based on category-level pose estimation [25,
26, 27, 28, 29]. On the other hand, [30] simplifies 6D pose into 3D pose,
i.e., the 2D object center and the orientation angle, and implements robotic
assembly in a plane. [31] enables robots to reason the 6D assembly pose by the
surrounding geometry and generalize to the objects that contain the shape
template. [32] further exploits 6D pose tracking to achieve assembly tasks
with tight tolerances, including box packing, cup stacking and marker
insertion. However, the 6D pose prediction above only relies on geometric
information; thus estimating 6D poses with language instructions is beyond
their reach.

###  II-B Robotic Manipulation with Language Instructions

The blossom of CV and NLP enables robots to comprehend natural language
instructions. Specifically, visual grounding tasks [3, 33, 34] predict the 2D
bounding boxes of the target objects referring to language descriptions. CLEVR
[35] contains a branch of cubes, cylinders and spheres with various attributes
and provides visual reasoning, including counting, comparison and logical
operations. Instead of asking questions about the scene as in CLEVR [35],
CLEVR-Ref+ [36] produces the bounding boxes and segmentation masks based on
the given referring expressions. As the blossom of transformer language models
[37, 38], the performance of visual grounding reaches a remarkable level [34].

Some works employ visual grounding methodology to locate objects for robotic
grasping. [4, 39, 5] propose to grasp the real-world objects by natural
language instructions to accomplish human-robot collaboration. [40] exploits
LSTM to directly detect grasps from textual command rather than localizing the
object, which avoids the error caused by inaccurate object localization. To
overcome the uncertainties of human language, [41] builds a partially
observable Markov decision process to track the history of observations, which
accomplishes interactively robotic grasping in clutter. However, the object is
merely localized by 2D bounding boxes or 2D grasp candidates. 2D location may
be sufficient for grasping, but high-precision robotic assembly exceeds its
capability.

Recently, [42] further combines large language models with reinforcement
learning to enable robots to comprehend long-horizon, abstract language
instructions. Considering the robot’s capability to follow the language
instructions given continuously, [6] enhances visual grounding tasks to
history-dependent ones. Additionally, [7] leverages 6D object poses to grasp
category-level objects, which is instructed by natural language. However, only
the 2D localization information is extracted within the text and the mapping
between natural language and 6D object pose is not established.

##  III Method

Figure 2: Overview of our system. Given an RGB image and a description text,
we first locate the base object and the target object by a 2D visual grounding
network MDETR [34]. The object ROIs are then fed into our proposed Language-
Instructed 6D Pose Regression Network (LanPose) to jointly predict the object
poses and the assembly pose. Specifically, our network predicts several
intermediate geometric feature maps and then directly regresses the 6D base
object pose. Simultaneously, the text encoding is produced by RoBRETa,
integrated along with geometric features by a cross-attention module and
mapped to the 6D assembly pose. Subsequently, an integrated robotic is
established to grasp, manipulate and assemble the target block precisely to
the base block.

###  III-A Problem Statement

Given an RGB image $\mathcal{I}$ and a natural language instruction
$\mathcal{W}$, we aim to estimate the 6D poses of the base object
$\mathcal{P}_{base}$, the target object $\mathcal{P}_{target}$ and the
assembly position $\mathcal{P}_{assembly}$ and then conduct robotic assembly.

Our framework is presented in Fig. 2. The RGB image $\mathcal{I}$ and the
language instruction $\mathcal{W}$ are fed to a 2D visual grounding network to
acquire the 2D bounding boxes of the base object $\mathcal{B}_{base}$ and the
target object $\mathcal{B}_{target}$.

$(\mathcal{B}_{base},~{}\mathcal{B}_{target})=\mathcal{F}_{1}(\mathcal{I},~{}\mathcal{W})$
(1)

Then the ROI of the base object along with the language instruction
$\mathcal{W}$ is fed to the Language-Instructed 6D Pose Regression Network to
jointly estimate the 6D poses of the base object $\mathcal{P}_{base}$ and the
assembly position $\mathcal{P}_{assembly}$.

$(\mathcal{P}_{base},~{}\mathcal{P}_{assembly})=\mathcal{F}_{2}(\mathcal{I},~{}\mathcal{B}_{base},~{}\mathcal{W})$
(2)

Meanwhile, the target object pose $\mathcal{P}_{target}$ is predicted through
the direct pose regression branch of the network.

$\mathcal{P}_{target}=\mathcal{F}_{2}(\mathcal{I},~{}\mathcal{B}_{target})$
(3)

Subsequently, the robotic system grasps the target object with the pose
$\mathcal{P}_{target}$ and assemble it to the base object with the pose
$\mathcal{P}_{assembly}$.

###  III-B Data Generation

We generate synthetic training and validation data based on BlenderProc [43],
a physically-based rendering methodology. To enhance the generalization
ability, the objects with arbitrary initial poses are dropped to planes with
different gradients, heights and material. The 2D bounding boxes and 6D poses
of objects are recorded by the physics engine Pybullet and the photorealistic
images are rendered via Blender, with the camera views sampled within a
certain range.

Subsequently, we follow [36] to generate referring expressions for the
rendered images, which correspond the object attributes and positions to the
language description. Several templates with different grammatical structures
(e.g., $\langle verb\rangle~{}\langle obj_{1}\rangle~{}\langle
prep\rangle~{}\langle obj_{2}\rangle$, $Grab~{}\langle obj_{1}\rangle$
$and~{}\langle verb\rangle~{}it$ $\langle prep\rangle~{}\langle
obj_{2}\rangle$) and the attributes of the objects including shape and color
are pre-defined. Additionally, we design 7 actions for robots to comprehend,
each with several pre-defined phrases (e.g., $insert~{}to$, $stack~{}on$,
$assemble~{}behind$), respectively. See Sec. IV-A for more details. The
corresponding 6D pose label of the assembly position is calculated by multiply
of the homogeneous transformation matrix of the base object and the relative
transformation between base object and assembly position, which is obtained by
Blender.

In general, the ground-truth 6D poses of the observed objects in the image are
required by the physics engine and the ground-truth 6D assembly poses are
jointly generated with the referring expressions, which differ according to
the desired actions. Aided by the photorealistic images along with augmented
text, our networks are trained solely with synthetic data. With the capability
of transferring directly to real scenes without further training, our
methodology is highly efficient and extensible, which avoids labor-consuming
real-world annotations.

###  III-C 2D Visual Grounding Network

In order to take advantage of the accuracy and runtime progress brought by the
rapidly growing 2D object detection community without altering or re-training
the pose network, most 6D object pose estimation methods [12, 10, 15]
disentangle 2D object detection and 6D object pose estimation process. In this
case, we adopt MDETR [34], an end-to-end modulated detector which detects
objects in an image instructed by a text query, to localize the target object
and the base object. Pre-trained on 1.3M text-image pairs, MDETR shows big
gains on phrase grounding, referring expression comprehension and segmentation
tasks. We subsequently fine-tune the network on our synthetic custom dataset,
using Adam optimizer with a learning rate of 5e-5.

To sum up, the RGB image captured by the camera is fed to MDETR along with the
text description to produce the 2D bounding boxes of the target object and the
base object.

###  III-D Language-Instructed 6D Pose Regression Network

After acquiring the 2D localization of the base object and the target object,
our network is required to integrate both image and text information to
predict 6D poses. For the target object, only the 6D pose itself needs to be
estimated for robotic grasping. However, for the base object, we would like to
predict the corresponding assembly position as well. We design Language-
Instructed 6D Pose Regression Network, namely LanPose, to simultaneously
predict the 6D poses of both base object and assembly position, as presented
in Fig. 2.

Direct 6D pose regression branch.  Given the ROI of the base object, we
directly regress its 6D pose. Specifically, several intermediate geometric
feature maps that indicate the 2D-3D correspondences are predicted to
supervise the network. Then a CNN module, which supersedes the traditional
RANSAC-based P$n$P algorithm, is leveraged to directly regress the 6D pose
$\mathcal{P}_{base}$. Noteworthy, the geometric feature maps also imply the
surrounding context of the target assembly position, which is crucial to the
6D pose prediction of the assembly position.

Language-integrated 6D pose mapping branch.  On the other hand, the text input
is fed into a pre-trained transformer language model RoBERTa [38] to yield a
sequence of hidden vectors, which is the same size as the text input. Then the
geometric feature maps, which contain the pose and shape context of the
observed object, and the text encoding, which indicates the relative
relationship of target assembly position towards the observed object, are
fused to produce an integrated image and text features. We exploit cross-
attention mechanism [44] to integrate image and language features, which
computes attention scores between the query, which in this case is the
geometric feature maps, and the key and value, which is the text encoding.
These attention scores are utilized to weight the values, which are then
aggregated to produce a joint representation of the integrated text-image
feature. Additionally, a CNN module is exploited to subsequently predict the
6D assembly position pose w.r.t. the camera $\mathcal{P}_{assembly}$.
Specifically, the CNN module consists of three convolutional layers, each
followed by Group Normalization and GELU activation. Two Fully Connected (FC)
layers are then applied to the global feature. Finally, two parallel FC layers
output the 3D rotation $\mathbf{R}$ and 3D translation $\mathbf{t}$,
respectively.

Losses.  Inspired by [10], we disentangle 6D assembly position pose loss as

$\mathcal{L}_{Assembly}=\mathcal{L}_{\mathbf{R}}+\mathcal{L}_{center}+\mathcal{L}_{z}.$
(4)

We employ point matching loss for rotation and parametrize rotation as the
first two columns of the rotation matrix following [15]. The translation is
further decoupled to the 2D center location and the distance w.r.t. the
camera. Therefore,

$\begin{cases}\mathcal{L}_{\mathbf{R}}&=\underset{\mathbf{x}\in\mathcal{M}}{\operatorname{avg}}\|\hat{\mathbf{R}}\mathbf{x}-\bar{\mathbf{R}}\mathbf{x}\|_{1}\\\
\mathcal{L}_{\text{center}}&=\|(\hat{\Delta}_{x}-\bar{\Delta}_{x},\hat{\Delta}_{y}-\bar{\Delta}_{y})\|_{1}\\\
\mathcal{L}_{z}&=\|\hat{\Delta}_{z}-\bar{\Delta}_{z}\|_{1}\end{cases},$ (5)

where $\mathcal{M}$ represents object model vertices, $\Delta$ indicates the
scale-invariant translation [10] and $\hat{\bullet}$ and $\bar{\bullet}$
severally denote estimation and ground-truth. For symmetric objects, we rotate
$\bar{\mathbf{R}}$ along the known symmetry axes to the best match of
$\hat{\mathbf{R}}$. The overall loss of LanPose can be summarized as

$\mathcal{L}_{LanPose}=\mathcal{L}_{Pose}+\mathcal{L}_{Geom}+\mathcal{L}_{Assembly},$
(6)

where $\mathcal{L}_{Pose}$ and $\mathcal{L}_{Geom}$ are defined in GDR-Net
[15].

Training strategy.  To enhance the precision and robustness of the network,
the direct 6D pose regression branch is required to solely reflect the
geometric context of the observed object without the effect of language
instructions. Thereby, we train direct 6D pose regression branch individually
in stage 1. Then the backbone, geometric heads and pose regression module are
frozen in stage 2 to provide surrounding features of the assembly position for
the language-integrated 6D pose mapping branch. Experimentally, we train the
network for 120 epochs in stage 1 and 100 in stage 2, using Ranger optimizer
with a batch size of 128 and a base learning rate of 1e-4.

##  IV Experiments

Figure 3: Visualization. (a) The 7 types of blocks we use. (b, c, d) Three
groups of language-instructed 6D pose regression results by rendering the 3D
models and overlaying the contours and coordinate axes on the image. Left: The
6D poses of the target object and the base object. Right: The 6D pose of
assembly position. Noteworthy, the blocks are placed on planes with various
gradients and heights.

###  IV-A Experimental Setup

We generate 39000 synthetic images for training and 1000 for validation (Fig.
3(b)). 7 types of blocks are employed for language-instructed 2D localization,
6D pose estimation and real-world robotic assembly experiments. Meanwhile, 7
assembly actions are designed, including $insert~{}to$, $combine~{}with$,
$stack~{}on$ and $assemble$ to the $front$, $right$, $left$ or $back$ of the
object, each with various descriptions (e.g., $combine~{}with$ &
$assemble~{}together~{}with$, $stack~{}on$ & $put~{}upon$). The corresponding
ground-truth assembly poses are designed to closely fit the base object to
conduct high-precision robotic assembly.

###  IV-B 2D Visual Grounding

Three backbones are employed for MDETR, as shown in Table. I. As can be seen,
0.794, 0.734 and 0.815 in Average Precision (AP) is derived with ResNet-101,
EfficientNet-B3 and EfficientNet-B5 backbones, respectively. Practically, we
exploit ResNet-101 backbone in our robotic assembly experiments as a trade-off
between accuracy and runtime. Two bounding boxes with the highest confidence
are selected for language-instructed 6D pose inference along with the
corresponding object categories.

TABLE I: Results for 2D Visual Grounding MDETR Backbone | AP | AP50 | AP75
---|---|---|---
ResNet-101 | 0.794 | 0.990 | 0.930
EfficientNet-B3 | 0.734 | 0.990 | 0.986
EfficientNet-B5 | 0.815 | 0.990 | 0.973

###  IV-C Language-Instructed 6D Pose Inference

TABLE II: Results for 6D pose estimation on validation set Block | ADD(-S) | 2∘ 2cm | 5∘ 5cm
---|---|---|---
0.02d | 0.05d | 0.1d
6D Object Pose
01 | 74.32 | 95.33 | 98.19 | 85.08 | 95.18
02 | 56.76 | 91.49 | 98.50 | 69.97 | 91.24
03 | 74.72 | 95.30 | 98.65 | 85.24 | 95.90
04 | 75.40 | 94.03 | 97.89 | 80.42 | 92.42
05 | 43.73 | 84.10 | 96.34 | 72.52 | 92.73
06 | 58.59 | 93.57 | 98.54 | 85.98 | 96.73
07 | 64.10 | 93.31 | 98.49 | 81.80 | 95.42
Mean | 63.95 | 92.45 | 98.09 | 80.14 | 94.23
6D Assembly Pose
01 | 57.24 | 91.85 | 98.09 | 45.67 | 92.56
02 | 27.40 | 67.40 | 88.00 | 22.00 | 74.60
03 | 46.64 | 93.08 | 99.40 | 55.87 | 93.18
04 | 32.90 | 75.63 | 91.88 | 24.17 | 78.84
05 | 20.90 | 54.37 | 83.02 | 24.42 | 84.12
06 | 34.27 | 80.30 | 96.18 | 40.70 | 93.17
07 | 37.84 | 84.58 | 98.30 | 56.66 | 95.30
Mean | 36.74 | 78.17 | 93.55 | 38.50 | 87.39

Evaluation metrics.  We exploit ADD(-S) and n∘ n cm to evaluate the 6D object
pose and 6D assembly pose inference performance. The ADD metric [45]
calculates the mean distance between the model points transformed by the
predicted pose and ground-truth pose. For symmetric objects, the ADD-S metric
[46] measures the mean distance from the closest points in two transformed
sets. The n∘ n cm measures whether the rotation error is less than n∘ and the
translation error is under n cm. The minimum error among all possible ground-
truth poses is taken for symmetric objects.

Analysis.  Table. II demonstrates the performance of our proposed LanPose. For
the direct 6D pose regression branch, we achieve 98.09 in ADD(-S)-0.1d and
94.23 in 5∘ 5 cm. With regards to more rigorous metrics ADD(-S)-0.05d and 2∘ 2
cm, 92.45 and 80.14 are derived, respectively. On the other hand, 93.55 in
ADD(-S)-0.1d and 87.39 in 5∘ 5 cm are attained for the assembly position pose
in Language-integrated 6D pose mapping branch. The lack of explicit geometric
features and the complexity of mapping the implicit geometric and linguistic
encoding to the 6D pose space prompt the difficulty in obtaining a competitive
performance with the observed object, yet it is sufficient for the robotic
assembly task. Additionally, the qualitative visualization results of real-
world data are shown in Fig. 3(c, d), which demonstrate the high precision of
our proposed network.

###  IV-D Inference Time Analysis

With an AMD Ryzen 7 5800 CPU and an NVIDIA RTX 3090 GPU, our approach takes
115 ms to integrate image and text input and produce 2D bounding boxes, and 36
ms for jointly predicting the 6D poses of the observed object and the
corresponding assembly position.

###  IV-E Robot Assembly Experiments

Figure 4: Visualization of robotic assembly experiments. (I) Success cases. (a) The 2D bounding boxes produced by visual grounding network. (b) The corresponding 6D pose estimation results of target object and assembly position. (c, d) The robotic grasping and assembly process, respectively. (II) Failure cases. (a) caused by 2D visual grounding. (b) caused by 6D pose estimation.  TABLE III: Success rate for robotic assembly experiments Block |  2D Visual Grounding |  Robotic Assembly
---|---|---
01 | 100.0% | 100.0% (20/20)
02 | 90.0% | 85.0% (17/20)
03 | 85.0% | 80.0% (16/20)
04 | 90.0% | 80.0% (16/20)
05 | 90.0% | 85.0% (17/20)
06 | 75.0% | 65.0% (13/20)
07 | 85.0% | 80.0% (16/20)
Mean | 87.9% | 82.1% (115/140)

We adopt KUKA LBR iiwa R820 robot arm and 2-finger Robotiq 2F-140 gripper to
establish the robotic assembly system, with an Intel RealSense D435i camera
mounted to the gripper. Robot experiments are executed by Robotic Operating
System (ROS), with iiwa stack [47] communicating between ROS and KUKA Sunrise
software.

Analysis.  We conduct 20 robotic assembly trials for each block. Specifically,
the blocks are placed with arbitrary poses in the cluttered background on the
plane with various gradients and heights, and the desired assembly position is
randomly given by natural language. As illustrated in Table. III, 87.9%
detection rate is achieved, i.e., the base and target objects are correctly
located by the 2D visual grounding network. We consider the trial with the
assembly error under 5 mm a successful one. Thus, 115 out of 140 trials are
successful, leading to a 82.1% total success rate of perception, grasping,
manipulation and assembly, as shown in Fig. 4(I). In real-world experiments,
various language expressions sharing the same purpose are employed and
repeatedly fed into the system, along with the arbitrarily placed blocks. Our
networks robustly locate the blocks and precisely predict the 6D assembly
pose, which is impregnable to the initial condition, proving that the
linguistic information is finely fused with geometric features and properly
mapped to the 6D pose space.

Failure cases.  The failures are provoked by two main reasons, the 2D visual
grounding failure and the 6D pose estimation error, as enumerated in Fig.
4(II). The objects in a cluttered background or stacked together with
occlusion may prompt detection mistakes. Accordingly, the inaccurate 2D
bounding boxes and pose ambiguity caused by occlusion or symmetry lead to 6D
pose estimation error. Practically, the 6D pose estimation precision required
by robotic grasping is not as strict as that demanded by robotic assembly.
Thus the robotic implementation failure regularly occurs in the assembly
process.

###  IV-F Ablation Study

TABLE IV: 6D assembly pose results of ablation study Method | ADD(-S) | 2∘ 2cm | 5∘ 5cm
---|---|---|---
0.02d | 0.05d | 0.1d
Ours | 36.74 | 78.17 | 93.55 | 38.50 | 87.39
w/o cross-attention | 21.93 | 63.94 | 88.75 | 15.09 | 65.61
$\mathcal{P}\rightarrow\Delta\mathcal{P}$ | 11.16 | 37.16 | 48.37 | 21.94 | 77.17

Ablation on cross-attention mechanism.  We compare the performance of
geometric feature fusion methodology, as presented in Table. IV. Instead of
exploiting the cross-attention mechanism to integrate image and language
features, we directly concatenate intermediate geometric feature maps and text
encoding in language-integrated 6D pose mapping branch. The experimental
results illustrate that the cross-attention mechanism surpasses the direct
concatenation of explicit geometric feature maps and plays a vital role on
modeling multi-modality relationships.

Ablation on pose type.  In the language-integrated 6D pose mapping branch, we
fuse text encoding and geometric feature maps to predict the absolute 6D
assembly pose w.r.t. the camera. As demonstrated in Table. IV, we compare the
results of predicting the relative 6D assembly pose w.r.t. the base object
with the baseline. Although the language description reflects the relative
position between the observed object and the desired assembly position, the
quantitative results of predicting $\Delta\mathcal{P}_{assembly}$ are not
competitive with directly predicting $\mathcal{P}_{assembly}$, which suggests
the explicit surrounding context provided by intermediate geometric maps
dominates the prediction of 6D assembly pose.

##  V Conclusion

In this paper, we establish an integrated robotic system to robustly assemble
blocks by natural language instructions. To this aim, Language-Instructed 6D
Pose Regression Network (LanPose) is proposed, which integrates image and
language information and maps them to the 6D pose in $SE(3)$ space. With the
two-stage training strategy, the 6D poses of the observed object and the
corresponding assembly position are jointly predicted without being coupled.
Meanwhile, the integration of intermediate geometric feature maps and text
encoding by the attention mechanism and the language-integrated 6D pose
mapping module enormously enhance the precision of the network.

Along with the reliable 2D visual grounding network, our system is capable of
perceiving, grasping, manipulating and assembling blocks robustly. Aided by
the photorealistic images generated by physically-based rendering and the
augmented text description, our networks trained solely with synthetic data
are able to transfer directly to real-world applications without further fine-
tuning. Our proposed methodology demonstrates the potential to more efficient
human-robot cooperation with natural language instructions.

## References

  * [1] S. Kazemzadeh, V. Ordonez, M. Matten, and T. Berg, “Referitgame: Referring to objects in photographs of natural scenes,” in _Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP)_ , 2014, pp. 787–798.
  * [2] J. Mao, J. Huang, A. Toshev, O. Camburu, A. L. Yuille, and K. Murphy, “Generation and comprehension of unambiguous object descriptions,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 11–20.
  * [3] L. Yu, P. Poirson, S. Yang, A. C. Berg, and T. L. Berg, “Modeling context in referring expressions,” in _European Conference on Computer Vision_. Springer, 2016, pp. 69–85.
  * [4] J. Hatori, Y. Kikuchi, S. Kobayashi, K. Takahashi, Y. Tsuboi, Y. Unno, W. Ko, and J. Tan, “Interactively picking real-world objects with unconstrained spoken language instructions,” in _2018 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2018, pp. 3774–3781.
  * [5] M. Shridhar and D. Hsu, “Interactive visual grounding of referring expressions for human-robot interaction,” _Robotics: Science and Systems (RSS)_ , 2018.
  * [6] H. Ahn, O. Kwon, K. Kim, J. Jeong, H. Jun, H. Lee, D. Lee, and S. Oh, “Visually grounding language instruction for history-dependent manipulation,” in _2022 International Conference on Robotics and Automation (ICRA)_. IEEE, 2022, pp. 675–682.
  * [7] C. Cheang, H. Lin, Y. Fu, and X. Xue, “Learning 6-dof object poses to grasp category-level objects by language instructions,” in _2022 International Conference on Robotics and Automation (ICRA)_ , 2022, pp. 8476–8482.
  * [8] M. Rad and V. Lepetit, “Bb8: A scalable, accurate, robust to partial occlusion method for predicting the 3d poses of challenging objects without using depth,” in _Proceedings of the IEEE international conference on computer vision_ , 2017, pp. 3828–3836.
  * [9] S. Peng, Y. Liu, Q. Huang, X. Zhou, and H. Bao, “Pvnet: Pixel-wise voting network for 6dof pose estimation,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2019, pp. 4561–4570.
  * [10] Z. Li, G. Wang, and X. Ji, “Cdpn: Coordinates-based disentangled pose network for real-time rgb-based 6-dof object pose estimation,” in _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2019, pp. 7678–7687.
  * [11] R. L. Haugaard and A. G. Buch, “Surfemb: Dense and continuous correspondence distributions for object pose estimation with learnt surface embeddings,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2022, pp. 6749–6758.
  * [12] Y. Xiang, T. Schmidt, V. Narayanan, and D. Fox, “Posecnn: A convolutional neural network for 6d object pose estimation in cluttered scenes,” in _Robotics: Science and Systems (RSS)_ , 2018.
  * [13] Y. Li, G. Wang, X. Ji, Y. Xiang, and D. Fox, “Deepim: Deep iterative matching for 6d pose estimation,” in _Proceedings of the European Conference on Computer Vision (ECCV)_ , 2018, pp. 683–698.
  * [14] Y. Labbé, J. Carpentier, M. Aubry, and J. Sivic, “Cosypose: Consistent multi-view multi-object 6d pose estimation,” in _European Conference on Computer Vision_. Springer, 2020, pp. 574–591.
  * [15] G. Wang, F. Manhardt, F. Tombari, and X. Ji, “Gdr-net: Geometry-guided direct regression network for monocular 6d object pose estimation,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2021, pp. 16 611–16 621.
  * [16] Y. Hu, P. Fua, W. Wang, and M. Salzmann, “Single-stage 6d object pose estimation,” in _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2020, pp. 2930–2939.
  * [17] Y. Di, F. Manhardt, G. Wang, X. Ji, N. Navab, and F. Tombari, “So-pose: Exploiting self-occlusion for direct 6d pose estimation,” in _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2021, pp. 12 396–12 405.
  * [18] A. Zeng, K.-T. Yu, S. Song, D. Suo, E. Walker, A. Rodriguez, and J. Xiao, “Multi-view self-supervised deep learning for 6d pose estimation in the amazon picking challenge,” in _2017 IEEE international conference on robotics and automation (ICRA)_. IEEE, 2017, pp. 1386–1383.
  * [19] J. Tremblay, T. To, B. Sundaralingam, Y. Xiang, D. Fox, and S. Birchfield, “Deep object pose estimation for semantic robotic grasping of household objects,” in _Conference on Robot Learning (CoRL)_ , 2018.
  * [20] G. Zhai, D. Huang, S.-C. Wu, H. Jung, Y. Di, F. Manhardt, F. Tombari, N. Navab, and B. Busam, “Monograspnet: 6-dof grasping with a single rgb image,” in _2023 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2023, pp. 1708–1714.
  * [21] G. Zhai, X. Cai, D. Huang, Y. Di, F. Manhardt, F. Tombari, N. Navab, and B. Busam, “Sg-bot: Object rearrangement via coarse-to-fine robotic imagination on scene graphs,” _arXiv preprint arXiv:2309.12188_ , 2023.
  * [22] X. Deng, Y. Xiang, A. Mousavian, C. Eppner, T. Bretl, and D. Fox, “Self-supervised 6d object pose estimation for robot manipulation,” in _2020 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2020, pp. 3665–3671.
  * [23] K. Chen, R. Cao, S. James, Y. Li, Y.-H. Liu, P. Abbeel, and Q. Dou, “Sim-to-real 6d object pose estimation via iterative self-training for robotic bin-picking,” _Proceedings of the European Conference on Computer Vision (ECCV)_ , 2022.
  * [24] B. Wen, W. Lian, K. Bekris, and S. Schaal, “Catgrasp: Learning category-level task-relevant grasping in clutter from simulation,” in _2022 International Conference on Robotics and Automation (ICRA)_. IEEE, 2022, pp. 6401–6408.
  * [25] Y. Di, R. Zhang, Z. Lou, F. Manhardt, X. Ji, N. Navab, and F. Tombari, “Gpv-pose: Category-level object pose estimation via geometry-guided point-wise voting,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2022, pp. 6781–6791.
  * [26] R. Zhang, Y. Di, Z. Lou, F. Manhardt, F. Tombari, and X. Ji, “Rbp-pose: Residual bounding box projection for category-level pose estimation,” in _European Conference on Computer Vision_. Springer, 2022, pp. 655–672.
  * [27] R. Zhang, Y. Di, F. Manhardt, F. Tombari, and X. Ji, “Ssp-pose: Symmetry-aware shape prior deformation for direct category-level object pose estimation,” in _2022 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_. IEEE, 2022, pp. 7452–7459.
  * [28] M. Zaccaria, F. Manhardt, Y. Di, F. Tombari, J. Aleotti, and M. Giorgini, “Self-supervised category-level 6d object pose estimation with optical flow consistency,” _IEEE Robotics and Automation Letters_ , vol. 8, no. 5, pp. 2510–2517, 2023.
  * [29] Y. Su, Y. Di, G. Zhai, F. Manhardt, J. Rambach, B. Busam, D. Stricker, and F. Tombari, “Opa-3d: Occlusion-aware pixel-wise aggregation for monocular 3d object detection,” _IEEE Robotics and Automation Letters_ , vol. 8, no. 3, pp. 1327–1334, 2023.
  * [30] Y. Litvak, A. Biess, and A. Bar-Hillel, “Learning pose estimation for high-precision robotic assembly using simulated depth images,” in _2019 International Conference on Robotics and Automation (ICRA)_. IEEE, 2019, pp. 3521–3527.
  * [31] S. Stevšić, S. Christen, and O. Hilliges, “Learning to assemble: Estimating 6d poses for robotic object-object manipulation,” _IEEE Robotics and Automation Letters_ , vol. 5, no. 2, pp. 1159–1166, 2020.
  * [32] A. S. Morgan, B. Wen, J. Liang, A. Boularias, A. M. Dollar, and K. Bekris, “Vision-driven Compliant Manipulation for Reliable, High-Precision Assembly Tasks,” _Robotics: Science and Systems XVII_ , 2021.
  * [33] L. Yu, Z. Lin, X. Shen, J. Yang, X. Lu, M. Bansal, and T. L. Berg, “Mattnet: Modular attention network for referring expression comprehension,” in _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , 2018, pp. 1307–1315.
  * [34] A. Kamath, M. Singh, Y. LeCun, G. Synnaeve, I. Misra, and N. Carion, “Mdetr-modulated detection for end-to-end multi-modal understanding,” in _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2021, pp. 1780–1790.
  * [35] J. Johnson, B. Hariharan, L. Van Der Maaten, L. Fei-Fei, C. Lawrence Zitnick, and R. Girshick, “Clevr: A diagnostic dataset for compositional language and elementary visual reasoning,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2017, pp. 2901–2910.
  * [36] R. Liu, C. Liu, Y. Bai, and A. L. Yuille, “Clevr-ref+: Diagnosing visual reasoning with referring expressions,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2019, pp. 4185–4194.
  * [37] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “Bert: Pre-training of deep bidirectional transformers for language understanding,” _arXiv preprint arXiv:1810.04805_ , 2018.
  * [38] Y. Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and V. Stoyanov, “Roberta: A robustly optimized bert pretraining approach,” _arXiv preprint arXiv:1907.11692_ , 2019.
  * [39] H. Ahn, S. Choi, N. Kim, G. Cha, and S. Oh, “Interactive text2pickup networks for natural language-based human–robot collaboration,” _IEEE Robotics and Automation Letters_ , vol. 3, no. 4, pp. 3308–3315, 2018.
  * [40] Y. Chen, R. Xu, Y. Lin, and P. A. Vela, “A joint network for grasp detection conditioned on natural language commands,” in _2021 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2021, pp. 4576–4582.
  * [41] H. Zhang, Y. Lu, C. Yu, D. Hsu, X. La, and N. Zheng, “Invigorate: Interactive visual grounding and grasping in clutter,” _Robotics: Science and Systems (RSS)_ , 2021.
  * [42] M. Ahn, A. Brohan, N. Brown, Y. Chebotar, O. Cortes, B. David, C. Finn, K. Gopalakrishnan, K. Hausman, A. Herzog _et al._ , “Do as i can, not as i say: Grounding language in robotic affordances,” _arXiv preprint arXiv:2204.01691_ , 2022.
  * [43] M. Denninger, M. Sundermeyer, D. Winkelbauer, D. Olefir, T. Hodan, Y. Zidan, M. Elbadrawy, M. Knauer, H. Katam, and A. Lodhi, “Blenderproc: Reducing the reality gap with photorealistic rendering,” in _International Conference on Robotics: Sciene and Systems, RSS 2020_ , 2020.
  * [44] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is all you need,” _Advances in neural information processing systems_ , vol. 30, 2017.
  * [45] S. Hinterstoisser, V. Lepetit, S. Ilic, S. Holzer, G. Bradski, K. Konolige, and N. Navab, “Model based training, detection and pose estimation of texture-less 3d objects in heavily cluttered scenes,” in _Asian conference on computer vision_. Springer, 2012, pp. 548–562.
  * [46] T. Hodaň, J. Matas, and Š. Obdržálek, “On evaluation of 6d object pose estimation,” in _European Conference on Computer Vision_. Springer, 2016, pp. 606–619.
  * [47] C. Hennersperger, B. Fuerst, S. Virga, O. Zettinig, B. Frisch, T. Neff, and N. Navab, “Towards mri-based autonomous robotic us acquisitions: a first feasibility study,” _IEEE transactions on medical imaging_ , vol. 36, no. 2, pp. 538–548, 2017.


# Tuna: A Static Analysis Approach to Optimizing Deep Neural Networks

1st Yao Wang  Amazon Web Services
<EMAIL_ADDRESS>2nd Xingyu Zhou  Amazon Web Services
<EMAIL_ADDRESS>3rd Yanming Wang  4th Rui Li  Amazon Web Services
<EMAIL_ADDRESS>University of Utah
<EMAIL_ADDRESS>5th Yong Wu  Amazon Web Services
<EMAIL_ADDRESS>6th Vin Sharma  Amazon Web Services
<EMAIL_ADDRESS>
###### Abstract

We introduce Tuna, a static analysis approach to optimizing deep neural
network programs. The optimization of tensor operations such as convolutions
and matrix multiplications is the key to improving the performance of deep
neural networks. Many deep learning model optimization mechanisms today use
dynamic analysis, which relies on experimental execution on a target device to
build a data-driven cost model of the program. The reliance on dynamic
profiling not only requires access to target hardware at compilation time but
also incurs significant cost in machine resources. We introduce an approach
that profiles the program by constructing features based on the target
hardware characteristics in order. We use static analysis of the relative
performance of tensor operations to optimize the deep learning program.
Experiments show that our approach can achieve up to $11\times$ performance
compared to dynamic profiling based methods with the same compilation time.

##  I Introduction

Deep neural networks (DNN) are the mainstay of many applications including
image classification, natural language processing, speech recognition, and
automated driving. Cloud service providers are offering machine learning
services that can compile, optimize, and deploy DNN on various target
hardware. Optimizing a large set of DNN models, from convolutional neural
networks to transformer-based networks, for a wide range of target hardware,
from general purpose processors to AI accelerator ASICs, is constrained by two
primary factors:

Compilation time and expense. A long compilation time can be a bad user
experience. In addition, the longer the compilation time, the greater the
infrastructure costs.

Cross compilation. A compiler cannot assume that it has access to the target
hardware and it must be able to cross-compile from a host machine with a
different hardware than the target.

Current machine learning compilers use two common ways to generate high
performance deep learning code for multiple target hardware. The first method
is based on auto-tuning, as implemented in AutoTVM [1], Ansor [2], FlexTensor
[3], and Halide auto-scheduler [4]. Systems using this approach usually
conduct a machine-learning driven auto-tuning search across a large predefined
program transformations space. The system generates code samples and measures
their performance on the target machine. It trains a machine learning model
based on the collected measurement data, and use this model to guide
subsequent code sample selection and code generation. The second method uses
vendor-supplied kernel libraries ($\mathit{e.g.}$, TensorRT, OneDNN). This is
widely adopted by current machine learning frameworks such as MXNet, PyTorch,
and TensorFlow. Both these methods, however, have several drawbacks that
restrict their application at scale. Although the auto-tuning approach
significantly reduces the engineering effort needed to hand-craft efficient
code for each target platform, training an effective cost model for that
target hardware requires collecting profile data from the execution of entire
DNN on a real device. This breaks the cross compilation constraints for
production compilation service. Furthermore, auto-tuning based compilation
requires large number of samplings in program transformation space to converge
to local optima, which can be prohibitively long. For example, tuning
TensorFlow SSD MobileNet for Amazon Graviton2 target with AutoTVM takes 240
hours. On the other hand, vendor kernel libraries provide efficient
handcrafted kernel codes which do not require auto-tuning. However, vendor
libraries don’t cover every operation used in DNN models. Adding more kernels
requires heavy engineering effort. In addition, one vendor library usually
only supports a specific hardware architecture. We have to integrate various
libraries into production service to support different target hardware, making
it hard to manage service infrastructure and extend to new target.

In order to resolve the challenge of building a compilation service with
limited compilation time and cross-compilation mechanism, this paper proposes
a system, Tuna, that utilizes the static analysis and a combination of
analytical cost modeling to optimizing DNN. Tuna has the following advantages
comparing to auto-tuning based compilation:

  * •

Optimization is done with static analysis and no real hardware is required.

  * •

Hardware-based cost model is transferable to different micro architectures.
Only one single cost model is required for one major architecture, such as CPU
and GPU.

  * •

Unlike performance measurement which requires sequential execution on a target
device, static analysis tasks can be fully paralleled on a multi-core CPU
machine, which largely reduces total analyzing time.

Comparing to vendor kernel libraries, Tuna supports a large set of DNN
kernels. Tuna provides a uniform compilation stack for various target
hardware.

In this paper, we create cost model for both CPU and Nvidia GPU architectures.
We evaluate our approach on cloud servers and edge devices. Experimental
results show our approach largely reduces total compilation time while
achieving better average performance comparing to dynamic profiled cost model.

##  II Overview

Tuna is a deep learning kernel optimization framework fully relying on
compile-time static analysis. Figure 1 shows the overall architecture of Tuna.
The system inputs are a tensor program $\mathit{e}$ together with
corresponding loop transformation candidate space $\mathit{T_{e}}$. For a
transformation $\mathit{t}$ $\in$ $\mathit{T_{e}}$, let $\mathit{i=g(e,t)}$ be
the intermediate representation of transformed program. Code generator then
generates low-level code $\mathit{a}$ ($\mathit{e.g.}$, assembly, PTX).
Hardware-based program feature is extracted via $\mathit{pf=f(i,a)}$. Program
performance score is calculated through $\mathit{c(pf)}$, with cost model
function $\mathit{c}$. The objective of Tuna is formalized as the following:

$\mathop{\arg\min}_{\mathit{t}\in\mathit{T_{e}}}\ \ \mathit{c(f(g(e,t),a))}.$
(1)

Tuna is built upon TVM [5], a deep learning compiler stack. While we reuse the
tensor expression system, IR and code generator from TVM, the following two
key components are the enhancements:

Hardware related cost model. Existing search-based approaches mainly rely on
high-level loop structure features, such as loop unrolling, loop vectorization
and number of arithmetic operations. We argue that these high level loop
features don’t include hardware specifications which have great impact on
kernel performance. Tuna analyzes both program IR and low-level generated
codes to extract target hardware related features, such as number of SIMD
instructions, CPU cache locality, GPU shared memory utilization, etc. With
these low-level hardware features, Tuna cost model is accurate enough to
predict the relative performance for a set of transformations of a tensor
program, rather than relying on experimental execution on real hardware
device.

Multi-threaded search algorithm. A multi-threaded search algorithm which can
fully utilize multi-core CPU resources is the key to accelerate searching. We
choose evolution strategies [6] as search algorithm. Evolution strategies is a
parallel black-box optimization algorithm. The computation among each
iteration of population can be executed simultaneously. With multi-core CPU
machine or a multi-machine compilation system, the total searching time can be
significantly reduced to be fit into the normal compilation flow.

Figure 1: Overview of the static analysis approach to optimizing tensor
programs.

##  III Hardware Related Cost Model

This section describes the analytical cost model to predict tensor program
performance. Tuna system can be formulated as following: we extract a series
of hardware related features, $\mathit{f_{0},f_{1},...,f_{n}}$, which have
great impact on deep learning tensor programs performance. These features can
fall into two categories:

Performance related instructions. Tuna counts the number of low-level
instructions that dominates tensor program performance, which mainly includes
arithmetic and data movement instructions. We propose a unified method which
jointly parses high-level program IR and low-level assembly code to get an
accurate estimation of instruction numbers.

General hardware features. Hardware features such as cache locality and
instruction level parallelism, have great impact on program performance.
Specific algorithms are required to extract such features from program IR and
assembly codes.

The program performance score is computed linearly with respect to features:

$\mathit{score=a_{0}*f_{0}+a_{1}*f1+...+a_{n}*f_{n}}$ (2)

The coefficients $\mathit{a_{0},a_{1},...,a_{n}}$ are generated for each
hardware architecture through hardware instruction latency and empirical
profiling data.

Tuna is generic enough to support different hardware architectures. This paper
covers Intel CPU, ARM CPU and Nvidia GPU. We also investigate the
transferability of the Tuna cost model across different micro architectures.
Experiment results of high-end sever level chips versus resource-limited edge
devices show that if two different micro architectures share the same set of
Single-Instruction-Multiple-Data (SIMD) related instructions, a single cost
model can be applied to both target platforms without modification.

###  III-A CPU cost model

This section describes Tuna CPU cost model in detail. To achieve decent
inference speed for deep learning program on modern CPU, it is common to do
the following optimizations: thread-level parallelism to utilize multi-core
CPU resources. Loop tiling to ensure good cache locality. Efficient
utilization of SIMD registers to maximize vector computation throughput.
Similarly, it’s necessary to include these hardware features as parts of our
CPU cost model to accurately analyze program performance. Tuna CPU cost model
includes the following features:

Number of SIMD instructions. Compute-intensive tensor program performance is
dominated by vector arithmetic and data movement instructions. For Intel AVX
instruction set, $\mathit{vfmadd}$ and $\mathit{vmov}$ are the most common
instructions in conv2d and dense operators, while for AARCH64 Neon
$\mathit{fmla}$, $\mathit{ld}$ and $\mathit{st}$ are used. We parse the
program IR and assembly codes to get the total number of these significant
SIMD instructions.

Estimation of L1 cache miss. Cache locality is a key factor for tensor program
on CPU. An analytical method is proposed in this paper to estimate L1 cache
miss for a given program IR.

Instruction level parallelism. Comparing to handcraft kernel libraries which
maximize fma instruction parallelism, search-based approach involves more
complicated instruction arrangements. It is significant to evaluate whether
the generated codes can keep CPU pipeline busy.

####  III-A1 Loop structure mapping

We started from counting the number of SIMD instructions. Program IR
represents the transformed tensor program in high level pattern, which
preserves the complete loop structures. However, the actual SIMD instruction
arrangement is opaque in high-level program IR, due to optimizations done in
code generation process, such as register allocation and superword-level
parallelism vectorization. On the other side, low-level assembly codes provide
detailed instruction information in each basic block. However, it is difficult
to restore original loop structure from assembly control-flow graph. It’s
infeasible to extract instruction information solely from either program IR or
assembly code.

In this paper we propose an algorithm which jointly parse high-level program
IR and low-level assembly code. The key idea is to extract full loop
information from program IR and low-level instruction quantity from assembly
code. A pattern matching is executed to match loop bodies with assembly basic
blocks and calculate total number of SIMD instructions.

Algorithm 1 shows the main procedure of jointly parsing. Loop blocks are
extracted from program abstract syntax tree with pre-order depth-first-search.
From assembly control flow graph, we identify a local basic block as a loop
candidate with the following condition: Traversing from top to bottom of
assembly code, there exists a jump instruction $\mathit{j}$ targeting a basic
block $\mathit{LBB_{x}}$, and the position of $\mathit{LBB_{x}}$ is above
$\mathit{j}$. Each pair of loop and basic block are then matched by checking
whether they have the same iteration boundary. For matched basic blocks, we
count the number of selected SIMD instructions. Finally, the algorithm
calculates the total number of SIMD instructions for every loop block in
original program IR.

Loop-Map(_$IR,assembly$_)

_ForLoops_ = Preorder-DFS-For-Loop(_$IR$_)

_LoopLBBs_ = Identify-Loop-LBB(_$assembly$_)

_matchedIdx = 0_

_matchedLBBs_ = []

__ for  _each basicBlock in LoopLBBs_ do

_forLoop_ = _ForLoops[matchedIdx]_

if  _Pattern -Match-Loop(_forLoop, basicBlock_)_ then

_matchedLBBs.append(basicBlock)_

_matchedIdx_ += 1

end if

end for

return _Count -Instruction(_$ForLoops,matchedLBBs$_)_

Algorithm 1 Algorithm for mapping program IR and assembly code

####  III-A2 Cache locality

In this section, we developed an improved data locality model as an analyisis
pass in TVM IR (TIR). This pass can analyze all kinds of computations
supported by TVM in a short time. Therefore it provides a fast data locality
approximation for the whole DNN network scheduling, and can be easily and
systematically combined with other important models and passes together to
guide the schedule selection of the whole network. The main idea of the model
is that we approximately estimate the data footprint and data movement volume
required to move into the cache of the object code. The object code was
abstracted as a tree consists of loop-nodes and access-nodes. All loop-nodes
are non-leaf nodes and carry all information of a loop statement. All access-
nodes are leaf nodes that represent either a tensor load access or a tensor
store access. The data footprint and data movement are calculated by
traversing the bottom to the top of the tree.

⬇

// Ni/Nj/Nk are perfect multiples of Ti/Tj/Tk

for(it = 0; it < Ni; it+=Ti)

for(jt = 0; jt < Nj; jt+=Tj)

for(k = 0; k < Nk; k++)

| for(i1 = 0; i1 < Ti; i1++)

| for(j1 = 0; j1 < Tj; j1++)

| C[i+it][j+jt]+=

| A[i+it][k]*B[k][j+jt];

for(l = 0; l < Nl; l++)

for(i2 = 0; i2 < Ti; i2++)

for(j2 = 0; j2 < Tj; j2++)

E[i+it][l]+=

C[i+it][j+jt]*D[j+jt][l]

Listing 1: Fused and tiled two matrix multiplication

We use Two Matrix Multiply (2MM) as an example to demonstrate the method for
building data movement. Listing 1 shows an example of fused and tiled 2MM
code. In this example the first Matmul uses $\mathit{i}$ and $\mathit{j}$ as
free index and perform contraction over $\mathit{k}$. The second Matmul uses
$\mathit{i}$ and $\mathit{l}$ as free index and perform contraction over
$\mathit{j}$. The tiling loop of $\mathit{i}$ and $\mathit{j}$ are tiled and
fused together, and other loops are non-tiled and non-fused. The cache
capacity and values of all tile-size and problem size are known before the
analysis starts.

We assume the cache capacity $\mathit{S}$ is enough to store all data
footprints below tiling loops, and is not enough to store data footprints of
any tiling loop. This means $\mathit{Ni}>\mathit{S}$ and
$\mathit{Nj}>\mathit{S}$, but
$\mathit{S}>\mathit{Ti}\mathit{Tj}+\mathit{Ti}\mathit{Nl}+\mathit{Tj}\mathit{Nl}+\mathit{Tj}\mathit{Nk}+\mathit{Ti}\mathit{Nk}$.

The data footprint and data movement are computed from leaf node to the root
of the tree. Since all leaf nodes are tensor accesses, the footprint and data
movement for leaf nodes are both $1$. The data footprint of a loop node is the
number of distinct data elements accessed during all iterations of this loop.
The data movement of a loop node is calculated based on its footprint and
cache capacity. If the footprint of a single loop iteration is smaller than
cache capacity, the data movement of this loop node is equal to its node
footprint. Otherwise, the data movement of this loop node is evaluated as
product of the number of iterations and the data movement of its single
iteration. The data movement of single iteration is correlated to data
movement of sub nodes, which is computed earlier in the bottom-up procedure.

In the 2MM example, based on our capacity assumption, the footprint of a
single iteration of loop $\mathit{jt}$ is
$\mathit{Ti}\mathit{Tj}+\mathit{Ti}\mathit{Nl}+\mathit{Tj}\mathit{Nl}+\mathit{Tj}\mathit{Nk}+\mathit{Ti}\mathit{Nk}$
which fits in cache. for all sub-loop nodes of $\mathit{jt}$, the data
movement is equal to data footprint. For loop $\mathit{jt}$, since its single
iteration footprint fit in cache, when the control flow goes to the next
iteration, tensor $\mathit{A}$ and $\mathit{E}$ could get reuse because their
access functions do not include index $\mathit{jt}$. Therefore, the data
movement of loop $\mathit{jt}$ is
$\mathit{Ti}\mathit{Nj}+\mathit{Ti}\mathit{Nl}+\mathit{Nj}\mathit{Nl}+\mathit{Nj}\mathit{Nk}+\mathit{Ti}\mathit{Nk}$,
which is still equal to its footprint. However, for loop $\mathit{it}$, even a
single iteration footprint does not fit in cache. So the data movement of loop
$\mathit{it}$ will be the product of number of iterations and data movement of
a single iteration, which is
$(\mathit{Ti}\mathit{Nj}+\mathit{Ti}\mathit{Nl}+\mathit{Nj}\mathit{Nl}+\mathit{Nj}\mathit{Nk}+\mathit{Ti}\mathit{Nk})*\mathit{Ni}/\mathit{Ti}$.
Figure 2 shows the loop structure and data movement calculation of the 2MM
example.

Algorithm 2 shows the main procedure of visiting each node in the TIR node
tree. When visiting a loop node, the algorithm recursively visits all its
children, and compute the union of data footprint of a single iteration. The
algorithm will detect whether the union of data footprint can fit in cache,
and track the reuse status of each tensor. If the data footprint of that
iteration fits in cache, the data movement volume is same as its footprint.
Otherwise, the algorithm will use the tracked reuse status information to
calculate the data movement. If the reuse status of the tensor is true, the
movement volume of that tensor will be equal to footprint, otherwise the
movement volume will be the movement volume of a single iteration multiply the
trip-count. The reuse status will be true at the leaf node by default. While
our analysis move from the bottom to top, the reuse status will be flipped to
false if the following conditions are true. The first case is if tensor
footprint exceed cache. The second case is if there exists a set of continuous
loop nodes that do not access this tensor, such that their footprints exceed
cache. Both implies the reuse distance of the tensor under discussion exceed
the cache capacity. The whole analysis module is implemented by using Integer
Set Library [7].

Visit-Node(_$node,cache$_)

if  _node is AccessNode_ then

return _Dmov=1, Dfp=1, DataSpace_

end if

else if  _node is LoopNode_ then

_subDspace, subDfp, subDmov_ = []

for  _each childNode_ do

_dspace_ , _dfp_ , _dmov_ = Visit-Node(_childNode, cache_)

append _dspace, dfp, dmov_ to _subDspace, subDfp, subDmov_

end for

 _DataSpace_ = Create-IntegerSet(_subDspace_)

if  _DataSpace.cardinality $>$ cache_ then

_Dmov_ = 0

for  _each tensor ts_ do

get _ts.dmov, ts.dspace, ts.expr_ from _subDmov, subDspace_

if  _ts.reuse is True_ then

_Dmov_ += Estimate-Dfp(_node.iterVar, ts.dspace_) Update-Reuse-Status(_ts_)

end if

else

_Dmov_ += _ts.dmov_ * _node.tripcount_

end if

end for

return _Dmov, Dfp=DataSpace.cardinality, DataSpace_

end if

else

return _Dmov=Dfp, Dfp=DataSpace.cardinality, DataSpace_

end if

end if

Algorithm 2 Algorithm for modeling data movement at each TIR tree node Figure
2: Loop structure and data movement of 2MM

####  III-A3 Instruction level parallelization

Modern CPU utilizes pipelines and out-of-order execution to increase the
instruction level parallelism (ILP) and hide the latency. Therefore, to
achieve high performance, the generated assembly code should allow the
processor issue as much instructions as possible to keep the pipeline busy.
Since the TVM framework includes more types of operators which will generates
more types of instructions, a specific model for specific operator is not
adequate to apply on the TVM-generated code. To solve this problem, we propose
a static analysis model to estimate the instruction level efficiency.

The key idea of the model is to design a simplified fast out-of-order
instruction scheduler that schedule instructions in each basic block. The
scheduler consists of two major components, the data dependency builder and
the instruction scheduler. The data dependency builder first scan the whole
basic block, and creates two instruction dependency graph for true dependency
and false dependency respectively. Then the instruction scheduler will
schedule the instructions based on the dependency graph and hardware
specifications such as instruction latency and number of different processing
unit. During the scheduling, a timestamp will be assigned to each instruction,
indicating the time point of the instruction starting executed. The first
ready-to-execute instruction will be scheduled at cycle zero, and total cycles
required for finalize all instructions will be used as the ILP cost of this
basic block. We calculate the product of ILP cost and number of executions for
a single basic block, and add up all products of all basic block. The
summation is used as the ILP cost of the whole program.

During the scheduling, the scheduler will manage two different hazard, the
structural hazard and data hazard. The structure hazard is controlled by
limiting the maximum number of instruction issued at each cycle. If the
maximum number of issued instruction reached the number of processing unit,
the next instruction to be issued will be delayed to next cycle. The data
hazard is identified by analyzing the dependency graph. If there is a read-
after-write (RAW) dependency between two instructions, the consumer
instruction should be scheduled after the producer instruction finishing
execution. If there is a write-after-read (WAR) dependency or write-after-
write dependency, the latter instruction who writes to the resource cannot be
scheduled before the prior instruction.

###  III-B Nvidia GPU cost model

This section describes our Nvidia GPU cost model in detail. To achieve
efficient execution, GPU tensor program needs to exhibit decent thread-level
parallelism to utilize GPU thread resources. Data locality is also significant
for shared memory efficiency. Our GPU cost model includes the following
features:

Number of PTX instructions. Similar to CPU assembly codes, we select
$\mathit{fma}$, $\mathit{ld}$ and $\mathit{st}$ as the most significant
instructions for tensor programs. We parse the Nvidia GPU PTX code to get the
total number of these instructions.

Thread level parallelism. Utilization of GPU threads largely determines the
performance of GPU kernels. We evaluate several aspects that directly affect
GPU thread-level parallelism: number of operations in a single thread,
Streaming Multiprocessor (SM) occupancy, warp latency hiding and shared memory
bank conflict.

####  III-B1 Loop structure mapping

The NVCC compilation process unrolls small loops with a known trip count by
default, which makes it hard to identify the corresponding loop structure in
high-level program IR from the low-level PTX code.

In this paper we propose an algorithm that parses the PTX code, identifies
loop structures and calculates total number of instructions. The key idea is
to identify the loop structure from PTX code and loop iterations from
registers.

Algorithm 3 shows the key idea of how we get the number of instructions. We
adopt the same idea from Algorithm 1 of identifying loop structures in
assembly code since PTX code has similar condition and jump structures. After
we have the loop structures, we maintain a register initial value map and
register value update map by parsing the PTX code. Since we already know the
loop structure, once we reach the line with eligible condition check, we know
it is for a certain loop. Based on the condition check, we get the register
used for comparing and the end condition. From the two maps we maintain, we
can easily calculate the loop iterations by initial value, update value and
end condition. Finally, the algorithm calculates the total number of PTX
instructions for every loop.

Loop-Map-PTX(_$PTX$_)

_LoopBBs_ = Identify-Loop-BB(_$PTX$_)

_regInitMap, regUpdateMap_ = Register-Match-Loop(__)

_LoopIterations_ = []

for  _each basicBlock in LoopBBs_ do

_loopIteration_ = Get-Iterations(_$regInitMap,regUpdateMap$_)

_LoopIterations.append(loopIteration)_

end for

return _Count -Instruction(_$LoopBBs,LoopIterations$_)_

Algorithm 3 Algorithm for identifying loop iterations in PTX code

####  III-B2 Thread level parallelization

In this section we discuss the details about features regarding to thread-
level parallelism.

Workload per thread We count the total number of PTX instructions for a single
thread with Algorithm 3. We calculate the total number of cycles based on
instruction cycles from [8]:

$\sum_{i=1}^{PTXInstruction}Count(i)*Cost(i)\ $ (3)

The workload per thread is represented by the total number of cycles.

Streaming Multiprocessor (SM) occupancy It is important to keep all SMs busy
for compute-intensive tensor programs. We check the total number of thread
blocks to determine whether it is greater than the total number of SMs, so
that all SMs have at least one block to execute. A penalty is added if the
total number of thread blocks is too small.

Warp latency hiding Warp scheduling is a vital mechanism to hide GPU
instruction latency. With more warps on each SM, GPU warp scheduler has a
better chance to hide memory latency, thus achieving better performance. We
calculate the maximum number of thread blocks that can be concurrently
scheduled in one SM by checking the number of registers and shared memory
usage per block. These information can be extracted with _nvcc ptxas-option_
command.

Shared memory bank conflict For modern Nvidia GPUs (compute capability $>=$
5.0), shared memory has 32 banks that are organized such that successive
32-bit words map to successive banks. Each bank has a bandwidth of 32 bits per
clock cycle and any shared memory load or store of n addresses that spans $n$
distinct memory banks can be served simultaneously, yielding an effective
bandwidth that is $n$ times as high as the bandwidth of a single bank.
However, if multiple addresses of a memory request map to the same memory
bank, the accesses are serialized. The hardware splits a memory request that
has bank conflicts into as many separate conflict-free requests as necessary,
decreasing the effective bandwidth by a factor equal to the number of separate
memory requests. The one exception here is when multiple threads in a warp
address the same shared memory location, resulting in a broadcast. In this
case, multiple broadcasts from different banks are coalesced into a single
cast from the requested shared memory locations to the threads. To incorporate
the effect of bank conflict, we first numerically evaluate the shared memory
access indices of all threads in the first warp from the IR to compute the
actual shared memory throughput. We then use the ratio between actual shared
memory throughput and requested shared memory throughput to adjust the number
of shared memory operations.

##  IV Search algorithm

In this paper, we need to search in a well-defined large search space to find
the optimal configuration. We choose Evolution Strategies (ES) as the search
algorithm and treat the search as an arbitrary black-box optimization problem.

ES works by treating the model of interest as an arbitrary optimization
problem. Given parameters of the model and the associated performance of that
model on the task of interest, ES is able to optimize and train the model
without any knowledge of the structure or architecture of the model itself.
Specifically, at each iteration, random gaussian noise is added to the
parameters to generate variations of the current model. Each variation is then
tested for performance on the task of interest. Finally, a weighted sum is
taken, based on performance, to generate an updated model. This becomes the
current model for the next iteration. In standard ES, we treat the performance
of the model within its environment as a black-box optimization problem. Now
we extend that concept to treat ES itself as a black-box optimization problem
with respect to the learning rate $\alpha$ and standard deviation of noise
$\sigma$.

Get learning rate $\alpha$, noise standard deviation $\sigma$, initial policy
parameters $\theta_{0}$

for  _t = 0, 1, 2, …_ do

Sample $\varepsilon_{1},...\varepsilon_{n}\sim\mathcal{N}(0,I)$

for  _i = 1, …, n_ do

$F_{i}=F(\theta_{t}+\sigma\varepsilon_{i})$

end for

$\theta_{t+1}=\theta_{t}+\alpha\frac{1}{n\sigma}\sum_{i=1}^{n}F_{i}\varepsilon_{i}$

end for

Algorithm 4 Evolution Strategies

The main steps about ES are given in Algorithm 4, $F$ denotes objective
function working with parameters $\theta$.

##  V Evaluation

This section evaluates the performance of Tuna. The following hardware
platforms are selected to cover the most widely used cases for deep learning
inference: Intel Xeon Platinum 8124M CPU (Amazon EC2 C5.9xlarge), AWS
Graviton2 (Amazon EC2 M6G.4xlarge, ARM 64-bit CPU), ARM Quad-core Cortex-A53
64-bit CPU (Acer aiSage), Nvidia Tesla V100 GPU (Amazon EC2 P3.2xlarge) and
Nvidia Jetson AGX Xavier GPU (512-core Volta GPU with Tensor Cores).

We evaluate Tuna on two levels: single operator and entire neural network. We
use AutoTVM as search-based tuning baseline for both levels of evaluation. We
also compare Tuna with state-of-the-art deep learning framework solutions for
the entire neural network evaluation.

We use TVM v0.7.0 release for all target devices. For Intel CPU targets, we
chose TensorFlow 1.15.3 and Pytorch 1.6.0 as deep learning framework
solutions. For AWS Graviton2 which has AARCH64 CPU, we use ARM Tool Solution
docker container which provides AARCH64 optimized version of TensorFlow and
Pytorch. For AWS P3 instance which has Nvidia V100 GPU, we use NGC docker
container with tag 21.03-tf1-py3 and pytorch:21.03-py3 which optimizes
TensorFlow and Pytorch for Nvidia GPU. For Nvidia Jetson AGX Xavier GPU, we
choose TensorFlow 1.15.2 and Pytorch 1.6.0.

###  V-A Entire Network

We first report Tuna performance on a set of deep learning models pre-trained
with TensorFlow and PyTorch, two popular deep learning frameworks: TensorFlow
SSD MobileNet v2, TensorFlow SSD Inception v2, PyTorch ResNet50 v1 and PyTorch
Bert base uncased. Selected models cover the most popular deep learning
applications: image classification, object detection and natural language
processing. Three metrics are measures to demonstrate the advantages of Tuna:

Compilation time. Table II records the compilation time of both Tuna and
AutoTVM to optimize neural networks. Experiments show that Tuna can speed up
compilation process up to $339\times$ comparing to AutoTVM.

Compilation cost. We measure the cost in dollar to optimize deep learning
models on Amazon EC2 instances to demonstrate the cost reduction provided by
Tuna. We multiply instance on-demand price with compilation time to get the
cost to compile. We benchmark Tuna on Amazon EC2 C5.24xlarge which takes $4.08
per hour. Table III shows Tuna reduces compilation cost down to 1.1% of the
original cost comparing to AutoTVM.

Inference latency. Table II shows the model inference latency from different
methods. The row of AutoTVM Partial shows the model inference latency compiled
with AutoTVM in the same compilation time of Tuna. Results show that with the
same compilation time, Tuna achieves up to $11\times$ performance of AutoTVM.
We also compare Tuna latency with AutoTVM full tuning models. On average Tuna
achieves 91.5% performance comparing to best possible schedules generated by
AutoTVM. This result shows that Tuna is able to achieve similar performance of
full tuning with significant less time and cost. Finally we compare Tuna
performance with deep learning framework solutions. Tuna achieves up to
$17.3\times$ performance. Due to memory resource limitation on Acer aiSage
device, directly running inference through deep learning framework is
infeasible. We only compare Tuna with AutoTVM on this platform.

###  V-B Single Operator

In this section we report the performance of a set of widely used compute-
intensive deep learning operators. We use AutoTVM tuned operator performance
as baseline. We reuse the optimization search space defined in AutoTVM to make
apple to apple comparison. We benchmark $\mathit{conv2d}$,
$\mathit{conv2d\\_winograd}$, $\mathit{depthwise\\_conv2d}$ and
$\mathit{batch\\_matrix\\_multiplication}$ performance on ARM CPU and Nvidia
GPU devices. We don’t measure $\mathit{conv2d\\_winograd}$ on Intel CPU device
since AutoTVM doesn’t define optimization space for this operator.

We define $\mathit{top-k}$ performance ratio as metric. Tuna generates top-k
best programs and execution latency of all programs are sum up. Similarly we
calculate the latency summation for top-k best programs generated with
AutoTVM. We divide two AutoTVM latency value with Tuna to evaluate the
effectiveness of Tuna to select decent optimizations from search space. Higher
value approaching 1 indicates Tuna can accurately predict the real execution
performance. Figure 3 and Figure 4 shows the benchmark result on Intel CPU,
ARM CPU and Nvidia GPU. On average We got 0.869 for top 10 and 0.873 for top
50. Experimental data shows Tuna is able to achieve quite close performance
comparing to full tuning method with AutoTVM.

Unit: ms | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
Framework | 22.09 | 24.29 | 31.26 | 253.04 |
AutoTVM Partial | 30.41 | 47.7 | 11.47 | 86.39 |
AutoTVM Full | 22.7 | 30.29 | 6.37 | 16.66 |
Tuna | 23.3 | 30.16 | 6.85 | 15.12 |

(a) Entire network performance on a system with Intel Xeon Platinum 8124M CPU

Unit: ms | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
Framework | 37.15 | 43.39 | 307.55 | 56.7 |
AutoTVM Partial | 46.95 | 75.37 | 195.38 | 63.62 |
AutoTVM Full | 29.55 | 45.56 | 16.11 | 15.11 |
Tuna | 30.24 | 44.3 | 17.77 | 16.13 |

(b) Entire network performance on a system with AWS Graviton2 ARM CPU

Unit: ms | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
AutoTVM Partial | 3499.72 | 4487.97 | 7111.88 | 1058.85 |
AutoTVM Full | 618.45 | 1419.83 | 1114.69 | 574.37 |
Tuna | 614.43 | 1389.35 | 1259.88 | 541.17 |

(c) Entire network performance on a system with ARM Quad-core Cortex-A53
64-bit CPU(Acer aiSage)

Unit: ms | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
Framework | 31.65 | 33.2 | 8.65 | 16.34 |
AutoTVM Partial | 44.77 | 70.97 | 105.16 | 13.03 |
AutoTVM Full | 27.43 | 30.69 | 2.34 | 3.24 |
Tuna | 28.45 | 34.87 | 2.65 | 4.62 |

(d) Entire network performance on a system with Nvidia V100 GPU

Unit: ms | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
Framework | 126.7 | 124.13 | 26.04 | 19.66 |
AutoTVM Partial | 202.21 | 426.44 | 21.16 | 51.32 |
AutoTVM Full | 50.78 | 57.25 | 15.21 | 7.64 |
Tuna | 59.81 | 80.05 | 18.82 | 11.49 |

(e) Entire network performance on a system with Nvidia Jetson AGX Xavier GPU

TABLE I: Entire network performance of _Tuna_ and the selected baselines.

Unit: hour | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
AutoTVM | 53 | 50 | 12 | 2.92 |
Tuna | 0.7 | 1 | 0.13 | 0.012 |

(a) Entire network compilation time for Intel Xeon Platinum 8124M CPU.

Unit: hour | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
AutoTVM | 240 | 280 | 49 | 3.2 |
Tuna | 3 | 5 | 0.45 | 0.062 |

(b) Entire network compilation time for AWS Graviton2 ARM CPU.

Unit: hour | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
AutoTVM | 299 | 316 | 64 | 5.14 |
Tuna | 3 | 5 | 0.45 | 0.062 |

(c) Entire network compilation time for ARM Quad-core Cortex-A53 64-bit
CPU(Acer aiSage).

Unit: hour | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
AutoTVM | 93.5 | 167.5 | 37 | 4.4 |
Tuna | 1.6 | 1.04 | 0.13 | 0.05 |

(d) Entire network compilation time for AWS P3 V100 GPU.

Unit: hour | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
AutoTVM | 202.5 | 356.3 | 83 | 4.4 |
Tuna | 1.7 | 1.05 | 0.13 | 0.05 |

(e) Entire network compilation time for Nvidia Jetson AGX Xavier GPU.

TABLE II: Entire network compilation time of _Tuna_ VS _AutoTVM_.

Unit: dollar | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
AutoTVM | 81.09 | 76.5 | 18.36 | 4.47 |
Tuna | 2.86 | 4.08 | 0.53 | 0.05 |

(a) Entire network compilation cost for Amazon EC2 C5.9xlarge(price $1.53 per
hour).

Unit: dollar | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
AutoTVM | 147.84 | 172.48 | 30.18 | 1.97 |
Tuna | 12.24 | 20.4 | 1.84 | 0.25 |

(b) Entire network compilation cost for AWS Graviton2 ARM CPU(price $0.616 per
hour).

Unit: dollar | TF SSD MobileNet | TF SSD Inception | PT ResNet50 | PT Bert |
---|---|---|---|---|---
AutoTVM | 286.32 | 512.61 | 113.22 | 13.46 |
Tuna | 6.53 | 4.24 | 0.53 | 0.2 |

(c) Entire network compilation cost for Amazon EC2 P3.2xlarge(price $3.06 per
hour).

TABLE III: Entire network compilation cost of _Tuna_ VS _AutoTVM_.

(a) Top-10 Performance Ratio on Intel Xeon Platinum 8124M CPU

(b) Top-10 Performance Ratio on AWS Graviton2 ARM CPU

(c) Top-10 Performance Ratio on AWS P3 V100 GPU

Figure 3: Top10 performance ratio for single operators from Tuna VS AutoTVM.

(a) Top-50 Performance Ratio on Intel Xeon Platinum 8124M CPU

(b) Top-50 Performance Ratio on AWS Graviton2 ARM CPU

(c) Top-50 Performance Ratio on AWS P3 V100 GPU

Figure 4: Top50 performance ratio for single operators from Tuna VS AutoTVM.

##  VI Related Works

Performance modeling has been a challenging problem due to the diversities of
hardware architectures. Over the years, researchers have proposed many
analytical approaches. Here we summarize notable related works.

CPU performance optimization: Several state-of-the-art prior works has focused
on analyzing the performance on CPU and generating efficient code. They
includes cache modeling tools [9][10], static tensor libraries[11][12][13],
model-driven tensor computation optimizers [14][15], polyhedral compilers and
polyhedral optimizations [16][17][18][19][20][21], and auto-tuning based
compilers [5][22][2]. However, all of the state-of-the-art researches cannot
be the off-the-shelf solution to our system for the following reasons. Cache
miss analyzers are accurate on analyzing the cache misses but are not able to
find the optimal loop schedule. Polyhedral compilers are usually separating
the tile-size selection and loop transformation, which makes its performance
not competitive to other frameworks. Auto-tuning based compilers are tuning
the loop schedule comprehensively for the whole network but results in very
long tuning and compiling time. Libraries are achieving high-performance on
some important operators but their static optimization may not be optimal for
all types of problem sizes, and they could block further graph-level
optimizations such as fusion. Model-driven tensor computation optimizers are
efficient and comprehensive, but they are not yet fully automated and can only
optimize specific computations such as tensor contractions and 2D
convolutions.

GPU performance modeling: Static analysis method has been used for GPU
performance modeling. Arafa et al. [8] developed a GPU performance prediction
framework named PPT-GPU, which utilizes a hybrid approach between analytical
modeling and cycle-accurate modeling. They used features such as instruction
statistics from the PTX code, instruction latencies, and L2 cache miss rate.
Guerreiro et al. [23] proposed a GPU performance evaluation method based on
recurrent neural networks and features extracted from PTX code. We need to
emphasize the goal of Tuna is to sort and find the best kernels in a defined
search space, while these approaches aim to accurately predict the performance
of different kernels. One significant advantage of Tuna is that it leverages
information from both low level assembly code and high level intermediate
representations.

Search-based auto-tuning: Search-based tuning approach has become an effective
method to optimize deep learning programs. AutoTVM [1] relies on user-
specified search space. Ansor [2] automatically generates search space for
small sub-graphs. Both AutoTVM and Ansor create machine-learning based cost
model and do training during exploration. Comparing to AutoTVM and Ansor, Tuna
applies hardware feature related cost model which doesn’t require any training
and accurately predict relative performance. Halide auto-scheduler [4]
generates schedules for the whole input program. All these search-based tuning
methods require experimental execution on target device, while Tuna fully
relies on static analysis.

##  VII Conclusion

In this paper, we proposed Tuna as an analytical approach to analyze and
optimize deep learning programs on modern CPU and GPU hardware. The
experiments show that we are able to achieve 99.21% average throughput for
three categories of deep learning models on various kinds of hardware
platforms, with merely 1.65% average compilation time comparing to search-
based tuning approach. Tuna achieves 1.54x average performance comparing to
state-of-the-art deep learning framework inference solutions. Similar
methodology can be applied to ASIC hardware. Extending Tuna to other
architecture is a future work.

## References

  * [1] Tianqi Chen, Lianmin Zheng, Eddie Yan, Ziheng Jiang, Thierry Moreau, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy.  Learning to optimize tensor programs.  arXiv preprint arXiv:1805.08166, 2018.
  * [2] Lianmin Zheng, Chengfan Jia, Minmin Sun, Zhao Wu, Cody Hao Yu, Ameer Haj-Ali, Yida Wang, Jun Yang, Danyang Zhuo, Koushik Sen, et al.  Ansor: Generating high-performance tensor programs for deep learning.  In 14th $\\{$USENIX$\\}$ Symposium on Operating Systems Design and Implementation ($\\{$OSDI$\\}$ 20), pages 863–879, 2020.
  * [3] Size Zheng, Yun Liang, Shuo Wang, Renze Chen, and Kaiwen Sheng.  Flextensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system.  In Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems, pages 859–873, 2020.
  * [4] Ravi Teja Mullapudi, Andrew Adams, Dillon Sharlet, Jonathan Ragan-Kelley, and Kayvon Fatahalian.  Automatically scheduling halide image processing pipelines.  ACM Transactions on Graphics (TOG), 35(4):1–11, 2016.
  * [5] Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, et al.  $\\{$TVM$\\}$: An automated end-to-end optimizing compiler for deep learning.  In 13th $\\{$USENIX$\\}$ Symposium on Operating Systems Design and Implementation ($\\{$OSDI$\\}$ 18), pages 578–594, 2018.
  * [6] Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever.  Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions.  arXiv preprint arXiv:1703.03864, 2017.
  * [7] Sven Verdoolaege.  isl: An integer set library for the polyhedral model.  In International Congress on Mathematical Software, pages 299–302. Springer, 2010.
  * [8] Yehia Arafa, Abdel-Hameed A Badawy, Gopinath Chennupati, Nandakishore Santhi, and Stephan Eidenbenz.  Ppt-gpu: Scalable gpu performance modeling.  IEEE Computer Architecture Letters, 18(1):55–58, 2019.
  * [9] Wenlei Bao, Sriram Krishnamoorthy, Louis-Noel Pouchet, and Ponnuswamy Sadayappan.  Analytical modeling of cache behavior for affine programs.  Proceedings of the ACM on Programming Languages, 2(POPL):1–26, 2017\.
  * [10] Tobias Gysi, Tobias Grosser, Laurin Brandner, and Torsten Hoefler.  A fast analytical model of fully associative caches.  In Proceedings of the 40th ACM SIGPLAN Conference on Programming Language Design and Implementation, pages 816–829, 2019.
  * [11] Tze Meng Low, Francisco D Igual, Tyler M Smith, and Enrique S Quintana-Orti.  Analytical modeling is enough for high-performance blis.  ACM Transactions on Mathematical Software (TOMS), 43(2):1–18, 2016\.
  * [12] Field G Van Zee and Robert A Van De Geijn.  Blis: A framework for rapidly instantiating blas functionality.  ACM Transactions on Mathematical Software (TOMS), 41(3):1–33, 2015\.
  * [13] Intel.  Intel oneapi deep neural network library (onednn).  https://software.intel.com/content/www/us/en/develop/documentation/oneapi-programming-guide/top/api-based-programming/intel-oneapi-deep-neural-network-library-onednn.html, 2020\.
  * [14] Rui Li, Aravind Sukumaran-Rajam, Richard Veras, Tze Meng Low, Fabrice Rastello, Atanas Rountev, and P Sadayappan.  Analytical cache modeling and tilesize optimization for tensor contractions.  In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–13, 2019.
  * [15] Rui Li, Yufan Xu, Aravind Sukumaran-Rajam, Atanas Rountev, and P Sadayappan.  Analytical characterization and design space exploration for optimization of cnns.  In Proceedings of the 26th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, pages 928–942, 2021.
  * [16] Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S Moses, Sven Verdoolaege, Andrew Adams, and Albert Cohen.  Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions.  arXiv preprint arXiv:1802.04730, 2018.
  * [17] Uday Bondhugula, Albert Hartono, Jagannathan Ramanujam, and Ponnuswamy Sadayappan.  A practical automatic polyhedral parallelizer and locality optimizer.  In Proceedings of the 29th ACM SIGPLAN Conference on Programming Language Design and Implementation, pages 101–113, 2008.
  * [18] Sven Verdoolaege, Juan Carlos Juega, Albert Cohen, Jose Ignacio Gomez, Christian Tenllado, and Francky Catthoor.  Polyhedral parallel code generation for cuda.  ACM Transactions on Architecture and Code Optimization (TACO), 9(4):1–23, 2013.
  * [19] Tobias Grosser, Hongbin Zheng, Raghesh Aloor, Andreas Simbürger, Armin Größlinger, and Louis-Noël Pouchet.  Polly-polyhedral optimization in llvm.  In Proceedings of the First International Workshop on Polyhedral Compilation Techniques (IMPACT), volume 2011, page 1, 2011.
  * [20] Martin Kong, Richard Veras, Kevin Stock, Franz Franchetti, Louis-Noël Pouchet, and Ponnuswamy Sadayappan.  When polyhedral transformations meet simd code generation.  In Proceedings of the 34th ACM SIGPLAN conference on Programming language design and implementation, pages 127–138, 2013.
  * [21] Cédric Bastoul.  Code generation in the polyhedral model is easier than you think.  In Proceedings. 13th International Conference on Parallel Architecture and Compilation Techniques, 2004. PACT 2004., pages 7–16. IEEE, 2004.
  * [22] Yizhi Liu, Yao Wang, Ruofei Yu, Mu Li, Vin Sharma, and Yida Wang.  Optimizing $\\{$CNN$\\}$ model inference on cpus.  In 2019 $\\{$USENIX$\\}$ Annual Technical Conference ($\\{$USENIX$\\}$$\\{$ATC$\\}$ 19), pages 1025–1040, 2019.
  * [23] João Guerreiro, Aleksandar Ilic, Nuno Roma, and Pedro Tomás.  Gpu static modeling using ptx and deep structured learning.  IEEE Access, 7:159150–159161, 2019.