ezyang / pytorch-unattached Goto Github PK

Weekend hacking = good time to do speculative, "not a good use of your time" thinking.

In #120 it was mentioned that we didn't have a convenient mechanism for formatting and printing user facing errors, and that we should somehow make it easier to do so. Continuing to make use of printf is probably the most practical way to proceed for now, but as I was working on the patch I felt bad for a few reasons:

1. printf is not type safe at all. Usually C programmers can just deal, but use of printf-style functions in C++ have their own hazards. The biggie: you have to remember NOT to pass std::string for a %s formatter.
2. printf style formatters make it difficult to render multiple line messages with increased indentation, which Haskell error messages use a lot, to good effect IMO. (Though, many would disagree with a claim that Haskell error messages are good.) While going down this rabbit hole, I noticed compilers like GCC/Clang have decided that multi-line error messages are not something you ever want to do (instead, you just emit extra diagnostic notes for any extra information you want to tack on.) This is probably precisely because their APIs don't make it convenient to output indented information...
3. printf style formatters discourage rendering of more complex data, like lists, etc., which may be of use to users.

So I turned my attention to the question: "What pretty-printing library could we use in C++ for this matter"? There are a few options:

1. @zdevito wrote a simple code templater in torch/csrc/jit/code_template.h which does type safe, template interpolation. Unfortunately, the API is a bit verbose as you have to explicitly build the env with explicit keys, specifying all types involved, and then apply it to the template. In principle, we could probably extend it to have a more compact API, but I don't know if this would destroy the conceptual integrity of the API.

2. Instead of printf, we could use C++'s existing printing convention using overload "<<" operators. Without any API, this is wordy, because you have to declare a string stream, write your error into it, and then extract out the string into your error message. But this can be solved, because temporaries in C++ live until the end of the full-expression in which they are called. So you can play tricks as in https://groups.google.com/forum/#!msg/comp.lang.c++/_GWLGQhbxYE/IDHpHMFm5XcJ to write the stringstream in one pipe. (Another version of the trick here: https://codereview.stackexchange.com/questions/6094/a-version-of-operator-that-returns-ostringstream-instead-of-ostream ). This API isn't going to give you indentation, but you don't have to write too much code.

3. I've always wanted an ACTUAL pretty-printer for C++, modeled off of, e.g., Wadler-Leijen pretty printers (https://homepages.inf.ed.ac.uk/wadler/papers/prettier/prettier.pdf). Depending on how much effort you put into it, there are two major jumps in functionality: first is (compositional) indentation (which comes from building a data structure representing the abstract document before rendering; much easier to re-indent now!); second is actually generating "pretty" output, whereby there are multiple layouts and the pretty-printer selects a layout that looks the best.

Personally, I'd like something that is (1) concise, (2) type safe and (3) which supports indentation, but I could be convinced that I don't actually want indentation.

I am not sure, since the C++ object has some stashed mutable state. Maybe it's OK because batchnorm is always used as a Module, and the Module is persistent across training iterations.

This means it cannot be used to case switch over const data.

Right now it assumes that it encapsulates the whole model, or that it's a unique subgraph in the model. If a cell is used multiple times, we'll record the first forward run and compile a one-stage closure when it's reused. This means that the closure isn't differentiable even once.

Hi @ezyang, I'd like to get my hands dirty with the IR code (great stuff BTW). I just built the jit branch and ran test_simple

x = Variable(torch.Tensor([0.4]), requires_grad=True)
y = Variable(torch.Tensor([0.7]), requires_grad=True)

torch._C._tracer_enter((x, y))
z = torch.sigmoid(torch.tanh(x * (x + y)))
trace = torch._C._tracer_exit((z,))

I get a segfault right at the last line. The segfault is due to THPGraphClass being nullptr in THPGraph_Wrap (in torch/csrc/jit/python_ir.cpp), in turn called at line 64 in torch/csrc/jit/python_tracer.cpp.

Apparently THPJIT_initExtension (in torch/csrc/jit/init.cpp) is never executed for me, so THPGraphClass is not initialized. Is there anything I'm missing on my end?

Thanks!

See here: https://github.com/ezyang/pytorch/blob/jit/test/expect/TestModels.test_word_language_model_LSTM.expect#L31

Needs investigation.

Our tests still crash on OSX because of the iterator stuff. We will need to fix this before we can release.

As long as we are in the business of inspecting Python operators (c.f. #36) to determine what they are, we need a more robust strategy for handling default arguments, e.g., as in this case:

class MaxPool1d(Function):
    @staticmethod
    def forward(ctx, input, kernel_size, stride=None, padding=0, dilation=1,
                ceil_mode=False):

The presence of default arguments in this method means that, when we see a MaxPool1d python operator, there may be as few as three arguments and as many as seven. Handling each of these cases can be quite annoying; if we had a mechanism for extracting the default arguments from the method definition (c.f., https://stackoverflow.com/questions/12627118/get-a-function-arguments-default-value) then we will have more robust code.

But this issue may be moot if we decide to stop inspecting Python operators, which is fragile anyway (since it could be changed from the Python level.)

If you somehow interact with Graph.op without having imported torch.toffee, the method will be missing (which will make you a sad panda). One pattern for handling this sort of thing is to have _C export a "base class" which is than extended by a Python class that adds helpers. The initialization code C code gets is hands on the Python class and then creates THAT, so the helpers are available. This is what is done for Variable.

Right now there is no public printer, and the << overload just prints the unique of the node. We need a more detailed printer.

@houseroad: This commit in ToffeeIR removes legacy attributes: https://github.com/ProjectToffee/ToffeeIR/blob/e1f50697ba376b0410d9e8b90d2607b735d0f1ec/toffee/frontend/c2_toffee.py

We need to update Pytorch so that it no longer generates these legacy attributes.

In particular we need to remove the attributes pad, kernel, dilation, stride, pad, adj, legacy_pad from Conv, ConvTranspose, MaxPool, and AveragePool. Instead use their plural equivalents.

You can get the number of dimensions of a node in a primspec function using: len(input.type().sizes())

Now that we have a fully dynamic Node interface, we no longer have exhaustive pattern matching on types of nodes. This means that we SHOULD NOT use IR_IF for cases where exhaustivity is important. We need to go through and audit these cases.

Fusion should be something that can be optionally applied after the fact. Deferring fusing will let us do end-to-end tests that don't depend on fusion working correctly. It also gives us a -O0 mode.

Right now all we get is the op name.

We have a helper that works for models, but nothing for plain old functions.

In particular, we have a public setType() function which has the effect of invalidating any Type* pointers returned from type(). This feels error-prone.

Right now, parameters are completely unnamed (just a big pile of numbers) because they are retrieved using parameters(). However, we actually have names for them from the state_dict interface from modules. It would be great to use these names to make the traces more interpretable.

We'll need to adjust name allocation to make use of name hints, and it's probably simplest if we unconditionally include the unique number even with the name hint.

Right now we only record sizes and strides. Knowing if it's actually a float/double/etc may also be useful.

Self explanatory

Right now the only uses we print are their corresponding Select nodes which isn't very useful

We recorded the types, we should print 'em.

When I trace AlexNet, the inplace ReLus turn into constants.

-  %18.0 = CppOp[ConvForward](%17, %1, %2), uses = [[]];
-  %20 = Constant(), uses = [%21.i0];
-  %21.0 = CppOp[ConvForward](%20, %3, %4), uses = [[]];
-  %23 = Constant(), uses = [%24.i0];
-  %24.0 = CppOp[ConvForward](%23, %5, %6), uses = [[]];
-  %26 = Constant(), uses = [%27.i0];
-  %27.0 = CppOp[ConvForward](%26, %7, %8), uses = [[]];
-  %29 = Constant(), uses = [%30.i0];

Figure out why!

We force it to be non-null and enforce it in the lint.

From yesterday's standup discussion.

Currently, the Select invariant is documented in text:

// Select nodes are used to handle multiple returns for the ops that actually return
// multiple values like PythonOp
// By convension, there is a unique select node for each output of an op
// so you can iterate over uses of a multi-return op to get all the select nodes.
// in this case
// number_of_outputs = op.uses().size()
// this will change if Tuples ever become first class.

This is not precise enough for me to write a machine check: the crux of the matter is, "What is a multi-return op?"

I propose that the multi-return-ness of a Node is statically determined by the kind of a node: in particular, with today's IR, only PythonOps are multi-return.

There are a few cases where they are defined in cpp files but that means they can't be used elsewhere.

We need to think carefully about the correctness conditions for optimizations involving inplace operations. Under the current semantics, when an variable is consumed by an inplace operation, PyTorch arranges so that there is never a reference to the old variable (because its tensor has been destroyed). We need to make sure that the optimizer never breaks this invariant, and we need a lint that can verify this is the case.

We also need to make sure reexecution of forwards respects data dependencies. At the moment it does not. See this test:

    def test_inplace_race(self):
        x = Variable(torch.FloatTensor([0]))
        @torch.jit.trace(num_derivatives=0)
        def fn(x):
            y = x + x + x + x + x + x + x + x
            x.add_(1)
            return x + y
        self.assertEqual(fn(x.clone()), y)
        self.assertEqual(fn(x.clone()), y) # run traced version

Adding to complexity is that data dependency can be masked by an aliasing operation, though I am having some difficulty coming up with a test for this case.

I saw this JIT_ASSERT while reviewing some code:

        // primspecs do not deal with Handles at the moment
        // so we map handles to Unused nodes
        auto typ = old->typeOption();
        JIT_ASSERTM(typ && typ->kind() == jit::TypeKind::HandleType,
          "primspec produced too few outputs");

This code is invoked when processing the outputs of a Python primspec. I don't think this should be an assert, because in my opinion, Python code counts as code "written by external users", and so of course they might misuse the API, and that is not an assert-failure (for which we should be permitted to compile without asserts in production), that is just a regular user failure.

There is one particularly awkward elephant in the room, however, which is that we are currently exposing the entire compiler IR API from Python for primspec construction, and it is dashingly easy to violate invariants from Python with this API. So, Python being the invariant violation barrier is not a hard rule--but I think failure to record enough outputs in a primspec squarely falls on the "raise an exception" side of the line.

CC @zdevito

Otherwise, you have to disable DataParallel modules to trace

We renamed it to MatchJITOps, so maybe graph_matcher.h/cpp is a better name, and it's uniform with graph_fuser. (Perhaps we should put these in a directory of their own.)

Depending on how many Evals we'll have in our traces in the end, and how large is a constant Eval overhead, we might want to implement an optimization pass that will stitch together subgraphs of multiple connected Evals so that they run a single larger node.

Right now it assumes opaque handles are never used. That will essentially only be true for forwards only graphs.

This code isn't exception safe:

template<pass_type optimizer>
PyObject * wrap_optimizer(PyObject *_unused, PyObject *py_graph) {
  HANDLE_TH_ERRORS
  THPUtils_assert(THPGraph_Check(py_graph), "expected a Graph instance");
  THPGraph *graph = (THPGraph*)py_graph;
  graph->cdata = optimizer(std::unique_ptr<Graph>{graph->cdata}).release();
  Py_RETURN_NONE;
  END_HANDLE_TH_ERRORS
}

If optimizer throws an exception, release() will never be called; so we'll deallocate graph->cdata, leaving the underlying Python object with a dangling pointer.

Doing a full run of the model test takes several minutes for me. This is bad, because it discourages people from running all of the tests. Is there anything we can do to make the tests run more quickly? A few possibilities: (1) run them on GPU, (2) have smaller parameters on the models.

Just because we've gone all dynamically typed on the IR, doesn't mean we can't have our symbols have a little more structure than just "strings". What we have found extremely useful in GHC is to be able to allocate one level of subnamespaces to our symbols. To keep things as strings, we achieve this by reserving the first letter of the symbol as the namespace signifier.

Looking at our list of built-in symbols, here are the namespaces I'd like to see:

• A Toffee Op namespace 'o' (oMul, oNegate)
• A Toffee Attribute namespace 'a' (aepsilon, amomentum)
• A "known key-attribute" namespace 'k' (kOffset, kInPlaceOutputs)

And you can always add more as necessary. The namespace signifier can be easily removed to get a raw string as necessary.

CC @zdevito

Right now, we have been adding operators to the AST one-by-one to support various operations we may be interested in. There are problems with doing this long term: every new operation you add, is another case that must be handled in any IR_IF case statement. IR_IF works well if you don't have too many forms in your language, not if you have a separate subkind for every single one of PyTorch's operations!

The object-oriented way to solve this problem is to keep subclasses, stop writing IR_IFs and move methods into the class definition. In GHC, we solved this problem differently, by maintaining an environment of "known" ops and recording with them information necessary to do analyses and optimizations with them (e.g., their types, strictness, unfoldings, etc.) I think this latter approach scales better with lots of ops (and is similar to how NNVM is structuring their operators: https://github.com/dmlc/nnvm/blob/master/docs/overview.md) but we should make a decision one way or another before we start adding tons and tons of operator definitions to fill out our PyTorch support.

Unlike NNVM, one thing I do NOT advocate is making this information public (at least for now.) So, no user-written ops that support fusion. That's a very tricky API to get right and there should be some concerted design that happens before we let people rely on it.

So that you can run caffe2 integration tests, but still have tests in PyTorch repo

We should at least to use a raw number, because Caffe2 sometimes prints them without a prefix.

https://github.com/ProjectToffee/ToffeeIR/issues/16

It would be really great if we could read out the debug information and then print it from the Python level. Would save a trip to gdb to get the trace when an exception is thrown.

Here is what I see when I run test/test_jit.py:

graph(%1, %2) {
  %3 = Add False(%1, %2);
  %4 =   %5 = Mul(%1, %3.0);
  %6 =   %7 = Tanh(%5.0);
  %8 =   %9 = Sigmoid(%7.0);
  %10 =   return (%9.0);
} 

For some reason, the output is being rendered twice, e.g. %4 = %5. Additionally, because we internally have select nodes, the output doesn't really match up with ToffeeIR's representation: we'll need some relabeling for select nodes to get the formats to match up.

(python3) [[email protected] ~/local/vision/torchvision/models] python alexnet.py                              
graph(%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17) {
  %18.0 = CppOp[ConvForward](%17, %1, %2), uses = [[]];
  %20.0 = CppOp[ConvForward](%21, %3, %4), uses = [[]];
  %21 = Constant(), uses = [%20.i0];
  %23.0 = CppOp[ConvForward](%24, %5, %6), uses = [[]];
  %24 = Constant(), uses = [%23.i0];
  %26.0 = CppOp[ConvForward](%27, %7, %8), uses = [[]];
  %27 = Constant(), uses = [%26.i0];
  %29.0 = CppOp[ConvForward](%30, %9, %10), uses = [[]];
  %30 = Constant(), uses = [%29.i0];
  %32.0 = ^Transpose(0, 1)(%11), uses = [[%35.i2]];
  %34 = Constant(), uses = [%35.i1];
  %35.0 = ^Addmm(1, 1, False)(%12, %34, %32.0), uses = [[]];
  %37.0 = ^Transpose(0, 1)(%13), uses = [[%40.i2]];
  %39 = Constant(), uses = [%40.i1];
  %40.0 = ^Addmm(1, 1, False)(%14, %39, %37.0), uses = [[]];
  %42.0 = ^Transpose(0, 1)(%15), uses = [[%45.i2]];
  %44 = Constant(), uses = [%45.i1];
  %45.0 = ^Addmm(1, 1, False)(%16, %44, %42.0), uses = [[%0.i0]];
  return (%45.0);
}

(python3) [[email protected] ~/local/vision/torchvision/models] python alexnet.py 
Traceback (most recent call last):
  File "alexnet.py", line 67, in <module>
    model(x)
  File "/data/users/ezyang/pytorch/torch/nn/modules/module.py", line 224, in __call__
    result = self.forward(*input, **kwargs)
  File "/data/users/ezyang/pytorch/torch/jit.py", line 137, in forward
    tuple(self.parameters()) + flatten(args))
  File "/data/users/ezyang/pytorch/torch/jit.py", line 22, in record_trace
    torch._C._jit_pass_lint(trace)
RuntimeError: torch/csrc/jit/ir.cpp:266: lint: Assertion `in_scope.count(input) == 1` failed.

I'm investigating

Because of the select invariant, the "proper" type of MultiType can always be reconstructed from the uses, but this is only if we never DCE select nodes. If we ever do DCE them, we may lose information about the type of the operator.

When I export a simple nn.Linear op, I get this:

version: 1
node {
  input: "1"
  output: "4"
  op_type: "Transpose"
  attribute {
    name: "axes"
    ints: 0
    ints: 1
  }
}
node {
  input: "3"
  input: "4"
  input: "2"
  output: "5"
  op_type: "FC"
}
name: "torch-jit-export"
input: "1"
input: "2"
input: "3"
output: "5"

The transpose is pretty pointless.

At the moment, we are using getPythonName to compute the correspondence between Python operands known wired-in ops. This is unsound for a few reasons. The most obvious is that __name__ only reports the unqualified name, which means that if someone else defines an Add function with different semantics, we'll incorrectly conclude that the operand can be used in this case. But the more insidious is that the C++ backend now has a dependency on the particular names being used in the Python level, when there shouldn't be any dependence at all.

An easy way to make things more robust is to look at the entire name, including module qualifier; then if we restrict ourselves to core library operands we can simply enforce as an invariant in PyTorch that people should not be renaming operands willy-nilly. You can avoid this invariant when we start moving operators to C++, in which case the C++ operand can be directly responsible for finding the IR. In the longest term, re #36, it would be best to not special-case on particular operator names, but design metadata for the operators which are sufficient for our passes to do what they need to do.

    def test_retrace_model(self):
        class MyNet(nn.Module):
            def __init__(self):
                super(MyNet, self).__init__()
                self.hidden = nn.Parameter(torch.randn(2, 2))
            def forward(self, input):
                return input
        m = MyNet()
        x = Variable(torch.randn(1))
        trace, _ = torch.jit.record_trace(m, x)
        y = Variable(torch.randn(1))
        trace, _ = torch.jit.record_trace(m, y)

Fails with:

======================================================================
ERROR: test_retrace_model (__main__.TestJit)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "test/test_jit.py", line 54, in test_retrace_model
    trace, _ = torch.jit.record_trace(m, y)
  File "/data/users/ezyang/pytorch/torch/jit.py", line 192, in record_trace
    args, parameters)
  File "/data/users/ezyang/pytorch/torch/jit.py", line 148, in record_trace
    self.saved_trace = torch._C._tracer_enter(trace_inputs)
RuntimeError: /data/users/ezyang/pytorch/torch/csrc/jit/tracer.h:167: enter: Assertion `input->tracing_state.state.expired()` failed.

I structured the example this way to avoid cries of "It's UB!" The code above is morally the same as this smaller repro:

    def test_retrace_model(self):
        def f(input):
            return input
        x = Variable(torch.randn(1))
        trace, _ = torch.jit.record_trace(f, x)
        trace, _ = torch.jit.record_trace(f, x)

And in case people may have forgotten, let me remind you of #57, which I claim would have solved this problem...

Because the hold a pointers to functions from arbitrary autograd graphs

While thinking about Adam's work on handling untraceable functions, which handles backwards to backwards-backwards, I realized that the analogous situation arises on forwards to backwards when we introduce a "stop tracing in this region" combinator. So it may make sense to design these in tandem. Here is a simple example:

@nojit
def f(y):
  return y * 3

y = x * 2
z = f(y)
a = z + 2
return a

This should trace into

graph (%x) {
  %y = MulScalar [2] %x
  %z, %f_grad_fn = CallPython [f] %y
  %a = AddScalar [2] %z
  return %a
}

With backwards tracing, this should trace into:

graph (%x, %grad_a) {
  %y = MulScalar [2] %x
  %z, %f_grad_fn = CallPython [f] %y
  %a = AddScalar [2] %z
  ----------- STAGE 1 -----------
  %grad_z = %grad_a
  %grad_y = Eval %f_grad_fn, %grad_z
  %grad_x = MulScalar [2] %grad_y 
  return %a, %grad_x
}

We have to handle the same edge-cases as in the backwards trace. For example, let's consider how to handle inplace updates.

@nojit_inplace(True)
def f(y):
  y.add_(2)
  return y * 3

y = x * 2
z = f(y)
a = z + y + 2
return a

Now when we trace this, we need to get the following:

graph (%x) {
  %y = MulScalar [2] %x
  %z, %y_inplace, %f_grad_fn = CallPython [f] %y
  %y_plus_two = AddScalar [2] %y_inplace
  %a = Add %z, %y_plus_two
  return %a
}

The key is the new %y_inplace output: in straight-line Python code, we can simulate inplace operations purely by adding them as extra outputs to the function, state monad style. Because %y becomes dead at this point, it is impossible for subsequent tracing to violate uniqueness. We can only do this if we know a priori which variables in the uninterpretable function might be mutated.

This issue says little about implementation issues atm, but here are some examples we want to support.

TODO: Do examples with multiple inputs and outputs.