Code

Functionalize in-place ops#584

Code

Merged

t-vi merged 38 commits intomainLightning-AI/lightning-thunder:mainfrom crpa/inplace-supportLightning-AI/lightning-thunder:crpa/inplace-supportCopy head branch name to clipboardJun 20, 2024

+714-52

Commits

Commits on Jun 20, 2024

• init: some binary ops

  Show description for c1a13fa
  crcrpar

  committed
  c1a13fa

  Copy full SHA for c1a13fa

• replace in-place with out-of-place

  Show description for 05afe9d
  crcrpar

  committed
  05afe9d

  Copy full SHA for 05afe9d

• correctly use prims.copy_ outputs in trace.output

  Show description for 351149a
  crcrpar

  committed
  351149a

  Copy full SHA for 351149a

• fix proxy swaps.

  Show description for c1d5234
  crcrpar

  committed
  c1d5234

  Copy full SHA for c1d5234

• preserve the change from upstream

  Show description for 0044f6d
  crcrpar

  committed
  0044f6d

  Copy full SHA for 0044f6d

• docstring

  Show description for 6355d34
  crcrpar

  committed
  6355d34

  Copy full SHA for 6355d34

• relax

  Show description for e348150
  crcrpar

  committed
  e348150

  Copy full SHA for e348150

• cover more in-place ops (#590)

  Show description for d7dca4d
  crcrpar

  committed
  d7dca4d

  Copy full SHA for d7dca4d

• store original and intermediate traces

  Show description for 6093ffe
  crcrpar

  committed
  6093ffe

  Copy full SHA for 6093ffe

• import cleanup

  Show description for 532abbb
  crcrpar

  committed
  532abbb

  Copy full SHA for 532abbb

• a bit simpler

  Show description for 82da20c
  crcrpar

  committed
  82da20c

  Copy full SHA for 82da20c

• guard copy if dst is an arg or a kwarg

  Show description for cf814f1
  crcrpar

  committed
  cf814f1

  Copy full SHA for cf814f1

• add OpTags.IN_PLACE

  Show description for ad6d84c
  crcrpar

  committed
  ad6d84c

  Copy full SHA for ad6d84c

• remove redundant checks

  Show description for acdab60
  crcrpar

  committed
  acdab60

  Copy full SHA for acdab60

• remove workaround for in-place add of num_batches_tracked

  Show description for b7ad1e2
  crcrpar

  committed
  b7ad1e2

  Copy full SHA for b7ad1e2

• clena up

  Show description for 203c396
  crcrpar

  committed
  203c396

  Copy full SHA for 203c396

• remove redundant copy bsym check

  Show description for 54f9fa4
  crcrpar

  committed
  54f9fa4

  Copy full SHA for 54f9fa4

• docstring

  Show description for 8a04a5e
  crcrpar

  committed
  8a04a5e

  Copy full SHA for 8a04a5e

• revert outdated diff

  Show description for e2a5c12
  crcrpar

  committed
  e2a5c12

  Copy full SHA for e2a5c12

• remove comment-out

  Show description for c7d13d0
  crcrpar

  committed
  c7d13d0

  Copy full SHA for c7d13d0

• is_inplace -> is_functionalizable

  Show description for 5afcbbe
  crcrpar

  committed
  5afcbbe

  Copy full SHA for 5afcbbe

• cleanup imports

  Show description for 367711f
  crcrpar

  committed
  367711f

  Copy full SHA for 367711f

• Update thunder/core/transform_common.py

  Show description for 1447071
  

  crcrpar

  and

  IvanYashchuk

  committed
  1447071

  Copy full SHA for 1447071

• maintain _inplace_to_out_of_place map and give in-place tag to the rest

  Show description for 8f6f23b
  crcrpar

  committed
  8f6f23b

  Copy full SHA for 8f6f23b

• support inplace args

  Show description for 08479ef
  crcrpar

  committed
  08479ef

  Copy full SHA for 08479ef

• register relu_

  Show description for 04808b5
  crcrpar

  committed
  04808b5

  Copy full SHA for 04808b5

• add inplace functionalization test

  Show description for 05605f7
  crcrpar

  committed
  05605f7

  Copy full SHA for 05605f7

• fix the bugs exposed by tests

  Show description for f88e6d4
  crcrpar

  committed
  f88e6d4

  Copy full SHA for f88e6d4

• remove debug prints

  Show description for ac5abe0
  crcrpar

  committed
  ac5abe0

  Copy full SHA for ac5abe0

• test ops with inplace args

  Show description for 494e08f
  crcrpar

  committed
  494e08f

  Copy full SHA for 494e08f

• comments

  Show description for efea02d
  crcrpar

  committed
  efea02d

  Copy full SHA for efea02d

• another comment

  Show description for 828ba29
  crcrpar

  committed
  828ba29

  Copy full SHA for 828ba29

• comments

  Show description for 2040d18
  crcrpar

  committed
  2040d18

  Copy full SHA for 2040d18

• relu_ calls relu(a, True)

  Show description for acd7655
  

  crcrpar

  and

  jjsjann123

  committed
  acd7655

  Copy full SHA for acd7655

• Revert "relu_ calls relu(a, True)"

  Show description for fa9756f
  crcrpar

  committed
  fa9756f

  Copy full SHA for fa9756f

• silu sig update

  Show description for bae091e
  crcrpar

  committed
  bae091e

  Copy full SHA for bae091e

• remove unused TYPE_CHECKING

  Show description for 3f03f74
  crcrpar

  committed
  3f03f74

  Copy full SHA for 3f03f74

• return t not return torch_func(...) for clarity

  Show description for d4402b0
  crcrpar

  committed
  d4402b0

  Copy full SHA for d4402b0