wschin / pytorch_compiler_tutorial

Codebase associated with the PyTorch compiler tutorial

This is the codebase associated with the PyTorch JIT compiler tutorial.

Build and Test

A recent version of PyTorch will likely need to have been installed (either nightly or built from source).

git clone https://github.com/bwasti/pytorch_compiler_tutorial.git --recursive
cd pytorch_compiler_tutorial
./build.sh
PYTHONPATH=build python test.py

This test benchmarks the following function with inputs of size 1024, run 100 times:

def foo(a, b):
  c = a.mul(b)
  a = c.mul(c)
  a = c.mul(a)
  return a

We expect to see this output:

-- Default IR --
 graph(%a.1 : Float(*),
      %b.1 : Float(*)):
  %c.1 : Float(*) = aten::mul(%a.1, %b.1) # test.py:20:7
  %a.3 : Float(*) = aten::mul(%c.1, %c.1) # test.py:21:7
  %a.5 : Float(*) = aten::mul(%c.1, %a.3) # test.py:22:7
  return (%a.5)

Default version took 26.74ms

-- Transformed IR --
 graph(%a.1 : Float(*),
      %b.1 : Float(*)):
  %a.5 : Float(*) = pw::CompilationGroup_0(%a.1, %b.1)
  return (%a.5)
with pw::CompilationGroup_0 = graph(%4 : Float(*),
      %5 : Float(*)):
  %c.1 : Float(*) = aten::mul(%4, %5) # test.py:34:7
  %a.3 : Float(*) = aten::mul(%c.1, %c.1) # test.py:35:7
  %a.5 : Float(*) = aten::mul(%c.1, %a.3) # test.py:36:7
  return (%a.5)

Compiled version took 8.20ms