pytorch_compiler_tutorial

KO | EN

gitlite — search

#python #hacktoberfest #go #golang #rust #linux #macos #nodejs #windows #awesome #cli #awesome-list #pytorch #mcp

pytorch_compiler_tutorial

(★ 47)

Codebase associated with the PyTorch compiler tutorial

pytorch_compiler_tutorial Latest Version Download

Download Latest Version (.zip)

// 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

pytorch_compiler_tutorial Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

pytorch_compiler_tutorial

pytorch_compiler_tutorial Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

Download README