AutoTVM & Device Fleet ` Learning to Optimize Tensor Programs - - PowerPoint PPT Presentation

AutoTVM & Device Fleet

`

Learning to Optimize Tensor Programs

High-level data flow graph and optimizations Hardware Frameworks

Learning to Optimize Tensor Programs

High-level data flow graph and optimizations Hardware Frameworks

Machine Learning based Program Optimizer

Learning to Optimize Tensor Programs

High-level data flow graph and optimizations Hardware Frameworks

Machine Learning based Program Optimizer

Learning to Optimize Tensor Programs

High-level data flow graph and optimizations Learning to generate optimized program for new operator workloads and hardware Hardware Frameworks

Search over Possible Program Transformations

Hardware

Loop Transformations Thread Bindings Cache Locality Thread Cooperation Tensorization Latency Hiding C = tvm.compute((m, n), 
 lambda y, x: tvm.sum(A[k, y] * B[k, x], axis=k))

Compute Description

Search over Possible Program Transformations

Hardware

Loop Transformations Thread Bindings Cache Locality Thread Cooperation Tensorization Latency Hiding C = tvm.compute((m, n), 
 lambda y, x: tvm.sum(A[k, y] * B[k, x], axis=k))

Compute Description

Search over Possible Program Transformations

Hardware

Loop Transformations Thread Bindings Cache Locality Thread Cooperation Tensorization Latency Hiding

Billions

• f possible
• ptimization

choices

C = tvm.compute((m, n), 
 lambda y, x: tvm.sum(A[k, y] * B[k, x], axis=k))

Compute Description

Learning-based Program Optimizer

4

Program Optimizer Program

Code Generator

Learning-based Program Optimizer

Runtime Measurements

4

Program Optimizer Program

Code Generator

Learning-based Program Optimizer

Runtime Measurements

High experiment cost, each trial costs ~1second

4

Program Optimizer Program

Code Generator

Learning-based Program Optimizer

5

Program Optimizer Program

Code Generator

Learning-based Program Optimizer

5

Program Optimizer Program

Code Generator

Cost Model

Learning-based Program Optimizer

Need reliable cost model per hardware

5

Program Optimizer Program

Code Generator

Cost Model

Learning-based Program Optimizer

Program Optimizer Program

Code Generator

Learning-based Program Optimizer

Program Optimizer Program

Code Generator

D

Training data

Learning-based Program Optimizer

Program Optimizer Program

Code Generator

D

Training data

Learning

Statistical Cost Model

Learning-based Program Optimizer

• Relatively low experiment cost
• Domain-specific problem structure
• Large quantity of similar tasks

Unique Problem Characteristics

Program Optimizer Program

Code Generator

D

Training data

Learning

Statistical Cost Model

Program-aware Cost Modeling

High-Level Configuration

Program-aware Cost Modeling

High-Level Configuration

for y in range(8): for x in range(8): C[y][x]=0 for k in range(8): C[y][x]+=A[k][y]*B[k][x]

Low-level Abstract Syntax Tree (shared between tasks)

Program-aware Cost Modeling

High-Level Configuration

for y in range(8): for x in range(8): C[y][x]=0 for k in range(8): C[y][x]+=A[k][y]*B[k][x]

Low-level Abstract Syntax Tree (shared between tasks)

C A B y 64 64 64 x 8 8 64 k 1 8 8

y 1 x 8 k 64

touched memory

• uter

loop length

statistical features Boosted Tree Ensembles

Program-aware Cost Modeling

High-Level Configuration

for y in range(8): for x in range(8): C[y][x]=0 for k in range(8): C[y][x]+=A[k][y]*B[k][x]

Low-level Abstract Syntax Tree (shared between tasks)

for context vec of x for for context vec of y context vec of k

+

soft scatter final embedding

TreeGRU

C A B y 64 64 64 x 8 8 64 k 1 8 8

y 1 x 8 k 64

touched memory

• uter

loop length

statistical features Boosted Tree Ensembles

Effectiveness of ML based Model

8

100 200 300 400 500 600 700 800 0.00 0.25 0.50 0.75 1.00 1.25 1.50

Effectiveness of ML based Model

8

100 200 300 400 500 600 700 800 0.00 0.25 0.50 0.75 1.00 1.25 1.50

One Conv2D Layer of ResNet18 on Titan X

Effectiveness of ML based Model

8

100 200 300 400 500 600 700 800 0.00 0.25 0.50 0.75 1.00 1.25 1.50

Number of Trials One Conv2D Layer of ResNet18 on Titan X

Effectiveness of ML based Model

8

100 200 300 400 500 600 700 800 0.00 0.25 0.50 0.75 1.00 1.25 1.50

Number of Trials One Conv2D Layer of ResNet18 on Titan X R e l a t i v e S p e e d u p Baseline: CuDNN

Effectiveness of ML based Model

9

100 200 300 400 500 600 700 800 0.00 0.25 0.50 0.75 1.00 1.25 1.50 100 200 300 400 500 600 700 800 0.00 0.25 0.50 0.75 1.00 1.25 1.50

Number of Trials One Conv2D Layer of ResNet18 on Titan X R e l a t i v e S p e e d u p Baseline: CuDNN TVM: Random Search

Effectiveness of ML based Model

10

100 200 300 400 500 600 700 800 0.00 0.25 0.50 0.75 1.00 1.25 1.50 100 200 300 400 500 600 700 800 0.00 0.25 0.50 0.75 1.00 1.25 1.50 100 200 300 400 500 600 700 800 0.00 0.25 0.50 0.75 1.00 1.25 1.50

One Conv2D Layer of ResNet18 on Titan X R e l a t i v e S p e e d u p TVM: Random Search TVM: ML-based Model Number of Trials Baseline: CuDNN

Transfer Learning Among Different Workloads

Historical Optimization Tasks

Transfer Learning Among Different Workloads

Historical Optimization Tasks Domain Invariant Program Representations

Transfer Learning Among Different Workloads

Historical Optimization Tasks Domain Invariant Program Representations Transferable Models to speedup new tasks

Transfer Learning Among Different Workloads

Historical Optimization Tasks Domain Invariant Program Representations Transferable Models to speedup new tasks

NVIDIA GPU Optimization (GTX 1080 Ti)

Latency (ms) 1.75 3.5 5.25 7 ResNet-50 MobileNet VGG-19 Inception V3 DenseNet-121

MXNet + TensorRT 4.0 AutoTVM

AMD GPU Optimization (Vega FE)

Latency (ms) 1.75 3.5 5.25 7 ResNet-50 MobileNet DenseNet-121

MIOpen AutoTVM

Bonus (INT8, GTX 1080)

Latency (ms) 0E+00 4E-05 8E-05 1.2E-04 1.6E-04 1-7-512-512-1 4-7-512-512-1 1-7-512-512-3 4-7-512-512-3 1-14-256-256-1 4-14-256-256-1

cuDNN AutoTVM

High Level: Scaling Automatic Performance Profiling

15

@ @ @ @

Fleet Tracker

Low Level: Portable RPC Tracker + Server

16

Resource Manager (Tracker) Nvidia GPU Server RPC RT CUDA tasks Android Phone RPC RT OpenCL tasks AMD GPU Server RPC RT ROCm tasks Zynq FPGA board RPC RT JIT driver Raspberry Pi RPC RT ARM tasks Shared cluster of heterogeneous devices Optimization Service RPC client Resource token Resource Allocation RPC Session Data Path Optimization Service RPC client cross compiler Red modules can be reconfigured remotely in each session cross compiler Running optimization services Prioritizer Workload 1 Workload 2 Workload 3 ML-based cost model … Hardware bitstream

RPC Communication Flow

17

Client Tracker Client Device upload code run code return run time Device

RPC Communication Flow

17

Client Tracker Client Device upload code run code return run time Device device free

RPC Communication Flow

17

Client Tracker Client Device upload code run code return run time Device device free

RPC Communication Flow

17

Client Tracker request device Client Device upload code run code return run time Device device free

RPC Communication Flow

17

Client Tracker request device Client Device upload code run code return run time Device device free

RPC Communication Flow

17

Client Tracker request device Client Device upload code run code return run time Device device free

RPC Communication Flow

17

Client Tracker request device return handle Client Device upload code run code return run time Device device free

Model to Tuned Implementation

18

Model Bag of Operators

• perator extraction

AutoTVM tuning

Tuned Model

Next: Autoscheduler, Lianmin @ 16:30

19

conv2d, x86 conv2d, GPU, winograd conv2d, ARM, spatial packing Handcrafted Schedule Templates