Associate fx node as origins on IRNode, ensure all nodes going to scheduling have an associate FX Node

[voznesenskym]

An example of what these IRNodes look like now, running test_alexnet_prefix_cpu

StorageBox(
  ComputedBuffer(name=None, layout=FlexibleLayout('cpu', torch.float32, size=[s3, 64, 55, 55], stride=[193600, 3025, 55, 1]), data=Pointwise(
    'cpu',
    torch.float32,
    relu(load(buf0, 193600*i0 + 3025*i1 + 55*i2 + i3, False) + load(primals_1, i1, False)),
    ranges=[s3, 64, 55, 55],
    origins={relu_default}
  ))
)
StorageBox(
  ComputedBuffer(name=None, layout=FlexibleLayout('cpu', torch.float32, size=[s3, 64, 27, 27], stride=[46656, 729, 27, 1]), data=Pointwise(
    'cpu',
    torch.float32,
    maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 112, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 111, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 110, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 57, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 56, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 55, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 2, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 1, False), load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3, False))))))))),
    ranges=[s3, 64, 27, 27],
    origins={max_pool2d_with_indices_default}
  ))
)
Convolution(
  name=buf0,
  layout=FixedLayout('cpu', torch.float32, size=[s3, 64, 55, 55], stride=[193600, 3025, 55, 1]),
  inputs=[InputBuffer(name='primals_3', layout=FixedLayout('cpu', torch.float32, size=[s3, s1, s4, s4], stride=[s1*s4**2, s4**2, s4, 1])), InputBuffer(name='primals_2', layout=FixedLayout('cpu', torch.float32, size=[s0, s1, s2, s2], stride=[s1*s2**2, s2**2, s2, 1]))],
  constant_args=(None, (4, 4), (2, 2), (1, 1), False, (0, 0), 1),
  output_view=None,
  origins={convolution_default}
)
ComputedBuffer(name='buf1', layout=FixedLayout('cpu', torch.float32, size=[s3, 64, 55, 55], stride=[193600, 3025, 55, 1]), data=Pointwise(
  'cpu',
  torch.float32,
  relu(load(buf0, 193600*i0 + 3025*i1 + 55*i2 + i3, False) + load(primals_1, i1, False)),
  ranges=[s3, 64, 55, 55],
  origins={relu_default}
))
ComputedBuffer(name='buf2', layout=FixedLayout('cpu', torch.float32, size=[s3, 64, 27, 27], stride=[46656, 729, 27, 1]), data=Pointwise(
  'cpu',
  torch.float32,
  maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 112, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 111, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 110, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 57, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 56, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 55, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 2, False), maximum(load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3 + 1, False), load(buf1, 193600*i0 + 3025*i1 + 110*i2 + 2*i3, False))))))))),
  ranges=[s3, 64, 27, 27],
  origins={max_pool2d_with_indices_default}
))

[jansel]

Could we just put everything on the IRNode base class? Having things scattered on specific node constructors doesn't seem great.

Something like:

class IRNode:
   _current_origins = set()
   @staticmethod
   @contexlib.contextmanager
   def current_origin(origins: Set[fx.Node])
      old = IRNode._current_origins
      IRNode._current_origins = old | origins
      yield
      IRNode._current_origins = old

   def __init__(self):
     self.origins = set(self._current_origins)

Then we could just throw a:

with IRNode.current_origins({n}):
   ...

Around the entire body of run_node(n).

[voznesenskym]

Yeah I thought about it, but I wasn't sure it was right for all IRNode to have this property on it. Is there any case where a subclass of IRNode having an origins is nonsensical?

[jansel]

No, every node can have an origins.

[jansel]

Unused?

[voznesenskym]

yes