[Papers] Learning to Optimize Tensor Programs

Low experiment cost, Domain-specific problem structure, Large quantity of similar operators.

Polyhedral models are a popular choice for S_e (search space); they model the loop domains as integer linear constrains. An alternative approach originating from halide defines a schedule space using a set of transformation primitives. –> This paper uses one in TVM to form S_e!

Learning to Optimize Tensor Programs

They use GBTs (XGBoost) and TreeGRU.

Training objective function. Common choice is the regression which encourages the model to predict cost accurately. They instead use rank loss function.

Use Simulated Annealing with the loss function as the energy function to explore. We select the top-performing batch of candidates to run on real hardware. Diversity-Aware Exploration is achieved by maximizing the objective to select candidate set S from top candidates (3). Uncertainty Estimator… I don’t know

Accelerating Optimization via Transfer Learning

In practice want to optimize for many tensor operators with different input shapes and data types. Apply transfer learning… Key is to create a transferable representation that is invariant to the source and target domains.

Common practice is to directly use configuration as the model’s input. However, search space specification can change for different workloads or when the user specifies a new search space for the same workload. The configuration representation is not invariant to changes in the search space.

Low-level loop AST is a shared representation of programs that is invariant to the search space. To leverage invariance, use low-level loop AST as input.

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s

This site uses Akismet to reduce spam. Learn how your comment data is processed.