Some refactoring #15

siddharth-krishna · 2020-12-23T17:44:56Z

This PR will make the following changes:

Rename Module to Function to be consistent with MLIR
Add a name property to Function
Unify type and shape inference into a single pass to be performed after the full Function is constructed
At op creation time, only look up op signature in terms of number of inputs and outputs
Use dataclasses to define core IR classes and use the frozen option to ensure they are immutable (should Functions be immutable? Shape/type inference happens after creation. Discussion needed.)
Removes any leftover uses of deepcopy.

So far, I've only done the renaming, but I'm opening the PR so we can discuss any issues early.

Also, this is going to touch pretty much every file, so we should figure out how to handle merge conflicts with the pipeline parallel PR.

santhnm2 · 2020-12-23T18:29:30Z

Those list of changes sound good, and good find on dataclasses - this is exactly the functionality I was trying to replicate in my commits yesterday re: removing ModuleView. :) Will this remove the need to do deep copies though?

dist_ir/ir/prettyprint.py

santhnm2 · 2020-12-29T22:45:45Z

I rebased this branch off of main after merging #12 - all the tests pass but there's an issue with prettyprint.py which I left a comment about

dist_ir/ir/op.py

* Make dataclass types hashable and update SequentialExecutor tests * Remove main from sequential executor tests * Fix pipeline parallel scheduler

siddharth-krishna

Thanks for your PR to this PR. :)
I had some minor questions which I've added as a review here since I merged your PR into this one.

dist_ir/ir/function.py

dist_ir/ir/op.py

dist_ir/executor/type_inference.py

dist_ir/executor/sequential_executor.py

dist_ir/ir/op.py

dist_ir/executor/type_inference.py

siddharth-krishna requested a review from santhnm2 December 23, 2020 17:45

siddharth-krishna mentioned this pull request Dec 24, 2020

Add pipeline parallel schedulers #12

Merged

Rename Module to Function

8dae1d6

santhnm2 force-pushed the refactor branch from bb40fad to 8dae1d6 Compare December 29, 2020 22:41

santhnm2 reviewed Dec 29, 2020

View reviewed changes

dist_ir/ir/prettyprint.py Outdated Show resolved Hide resolved

siddharth-krishna mentioned this pull request Jan 4, 2021

Add MLIR parser #14

Merged

siddharth-krishna added 5 commits January 4, 2021 14:50

Add device_var attribute to Pmap op

92f6724

Fix prettyprinting of Pmap

a3bdb0f

Use frozen dataclasses for Types

b4f361d

Convert Device to frozen dataclass

b7cc2c1

Convert Op to a frozen dataclass

91e4022

santhnm2 reviewed Jan 4, 2021

View reviewed changes

dist_ir/ir/op.py Outdated Show resolved Hide resolved

siddharth-krishna and others added 3 commits January 5, 2021 16:04

Convert Function to frozen dataclass; add FunctionMaker

be7ea6b

Convert Value to frozen dataclass

98bd6b6

Make dataclass types hashable and use Value instead of value names (#17)

dd215e0

* Make dataclass types hashable and update SequentialExecutor tests * Remove main from sequential executor tests * Fix pipeline parallel scheduler

siddharth-krishna commented Jan 6, 2021

View reviewed changes

dist_ir/ir/function.py Show resolved Hide resolved

dist_ir/ir/function.py Show resolved Hide resolved

dist_ir/ir/function.py Outdated Show resolved Hide resolved

dist_ir/ir/op.py Outdated Show resolved Hide resolved

siddharth-krishna and others added 12 commits January 6, 2021 17:04

Upgrade Python to 3.8, cache pip

50b71e7

Op: minor simplification

7521683

Merge branch 'main' into refactor

fea8544

Update MLIR parser to use refactored IR

8f398e4

Allow passing output types to Op()

fa22c1b

Fix typo, import order

afbf893

First implementation of type inference

c793996

Fix a test

cbdfd28

Add some TODOs

8c1b280

Replace strings with values in data parallel transform (and tests)

b911d27

Add remaining type inference functions

eaa43f9

Add type inference for pmap

84145d3

Change names to values for pipeline parallel transform

84f7a99

siddharth-krishna commented Jan 12, 2021

View reviewed changes

dist_ir/ir/op.py Outdated Show resolved Hide resolved

siddharth-krishna commented Jan 12, 2021

View reviewed changes

dist_ir/executor/type_inference.py Outdated Show resolved Hide resolved

siddharth-krishna added 5 commits January 12, 2021 18:02

Implement pmap's type prop fn, fix some others

f83cd3a

Fix Device constructor arg order

96538d6

Misc

464e9a6

Fix tests

9ea9c5d

Move shape inference to type inference

39d4f06

siddharth-krishna commented Jan 12, 2021

View reviewed changes

dist_ir/executor/sequential_executor.py Outdated Show resolved Hide resolved

siddharth-krishna added 2 commits January 12, 2021 18:14

Remove shape inference module

c5448f0

Remove uses of deepcopy

b5f8062

siddharth-krishna commented Jan 12, 2021

View reviewed changes

dist_ir/ir/op.py Outdated Show resolved Hide resolved

santhnm2 reviewed Jan 12, 2021

View reviewed changes

dist_ir/executor/type_inference.py Show resolved Hide resolved

siddharth-krishna and others added 7 commits January 12, 2021 18:38

Rename Op.{in,out}_edges to Op.{in,out}puts

c6729f1

Use ops instead of names for pipeline parallel utils

ac27d8b

Fix for error message when output types and outputs don't match

8b69a30

Remove TODO for replacing value names with values

dfb6168

Make type inference functions more robust

10caebf

Clean up op register

8d51621

Minor improvements to type prop fns

bfec0c3

siddharth-krishna merged commit 3c9905b into main Jan 12, 2021

siddharth-krishna deleted the refactor branch January 12, 2021 22:16

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Some refactoring #15

Some refactoring #15

siddharth-krishna commented Dec 23, 2020

santhnm2 commented Dec 23, 2020

santhnm2 commented Dec 29, 2020

siddharth-krishna left a comment

Some refactoring #15

Some refactoring #15

Conversation

siddharth-krishna commented Dec 23, 2020

santhnm2 commented Dec 23, 2020

santhnm2 commented Dec 29, 2020

siddharth-krishna left a comment

Choose a reason for hiding this comment