pytorch_compiler_tutorial

KO | EN

gitlite — search

#php #java #javascript #ruby #c-plus-plus #python #r #docker #android #rails #c-sharp #arduino

pytorch_compiler_tutorial

★ 47 Open GitHub ↗

Codebase associated with the PyTorch compiler tutorial

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

This is the codebase associated with the [PyTorch JIT compiler tutorial](https://jott.live/markdown/Writing%20a%20Toy%20Backend%20Compiler%20for%20PyTorch).

## 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
```

pytorch_compiler_tutorial

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README