Reduce the work to add a new expr: unify the structure of exprs #2190

zasdfgbnm · 2022-11-15T23:54:40Z

I added a new data member to Expr: std::vector<Val*> attributes_. Subclasses of Expr are now not allowed to have their own data member in the class. Instead, if a data member is neither an input nor an output, then the subclass should register it as an attribute. Some subclass of Expr needs to store plain data that is not a Val, such as DataType, bool, etc. These plain data are not Val, therefore impossible to store as an attribute. To work around this, I created a class template <typename T> class Attribute : public Val which can be used to store data of type T.

By following the contract that subclass of Expr must store data inside inputs, outputs, or attributes, every expr type can be constructed with the following ctor:

Expr(IrBuilderPasskey, std::vector<Val*> inputs, std::vector<Val*> outputs, std::vector<Val*> attributes);

And IR cloning and shallow copying now become independent of expr type. Now adding a new expr type does not need to manually write down the cloning constructor, sameAs, and shallowCopy. They can all be inherited from Expr. This is a big save that helps me to reduce ~800 lines of code.

I also refactored IrCloner to move the type dispatch as a member function.

I believe that with my change in this PR, we can also greatly simplify IR mutator. But I will not do it in this PR.

I recommend starting the review from torch/csrc/jit/codegen/cuda/ir_base_nodes.h.

…expr

This reverts commit c0a5864.

…ctured-expr

zasdfgbnm · 2022-11-17T09:09:23Z

torch/csrc/jit/codegen/cuda/ir_builder.cpp

@@ -8,74 +8,6 @@ namespace jit {
 namespace fuser {
 namespace cuda {

-//! Clone an IR node, forwarding the arguments to the IrCloner constructor.


moved to header file

zasdfgbnm · 2022-11-17T09:10:01Z

torch/csrc/jit/codegen/cuda/ir_builder.h

@@ -15,20 +16,6 @@ class Kernel;

 class IrCloner;

-// Passkey for builder to register properties with statements, and to call


Moved to a separate file torch/csrc/jit/codegen/cuda/ir_builder_key.h so it can be included in torch/csrc/jit/codegen/cuda/ir_base_nodes.h.

Oh, I did that too in #2197 😄

zasdfgbnm · 2022-11-17T10:05:11Z

torch/csrc/jit/codegen/cuda/ir_base_nodes.h

  // Creates a new instance of the expression with all its field copied.
  // Note that unlike IrCloner, this function only do a shallow copy
-  virtual Expr* shallowCopy() const = 0;
+  Expr* shallowCopy() const;


This method now is not even virtual

naoyam · 2022-11-17T21:57:01Z

torch/csrc/jit/codegen/cuda/ir_base_nodes.h

+#define DECLARE_CLONE \
+  virtual Statement* clone(IrCloner* ir_cloner) const override;
+
+#define DEFINE_CLONE(ClassName)                            \
+  Statement* ClassName::clone(IrCloner* ir_cloner) const { \
+    return IrBuilder::clone(this, ir_cloner);              \
+  }


Can we undef these somewhere? Otherwise, we should use names like PYTORCH_NVFUSER_DECLARE_CLONE.

I don't think so. These macros will be used in cpp files. I will rename it as NVFUSER_DECLARE_CLONE.

naoyam · 2022-11-17T21:58:10Z

torch/csrc/jit/codegen/cuda/ir_base_nodes.h

+//! to hold non-IR data, such as DataType, std::vector<int>, etc. Please don't
+//! use this class to hold IR nodes or their pointers.
+template <typename T>
+class TORCH_CUDA_CU_API PlainVal : public Val {


So, this is intended to hold non-IR data, but then why is it a subclass of Val?

So, this is used to hold IR node attributes. Does it need to be a subclass of Val?

Because attributes_ are vectors of Val* and I want to push plain values into attributes_. Maybe I should change attributes_ to a vector of Statement*? Currently, lots of kernel IR has expr members, and I am storing them as something like PlainVal<Statement*>; maybe making attributes_ a vector of Statement* would make it easier because I can directly store it.

Ah, I see. Yeah, we want to store proper IR-node values as well. Can it be std:vector<std::any>?

Not today. std::any is C++17 but PyTorch now is still using C++14

Ah, right. Can it be a vector of unique_ptr?

I just feel odd to have arbitrary attributes, like a DataType and a bool flag, as a Val. Please convince me if you really want to use PlainVal.

I believe the biggest advantage of making it a Val is that, IrCloner will automatically help you do deep copy when necessary. For example, if I have an attribute bool, and I clone the expr, then I will also have a new PlainVal<bool> object for the new expr. If later, I changed the new expr's bool attribute, the old one's bool attribute is not changed. But if we manage this by hand, we could make mistakes that, when cloning an expr, the old and new expr are pointing to the same object for its attribute, and if we change one, the other is also changed.

So to answer your question:

Can it be a vector of unique_ptr?

Yes. Actually, it SHOULD be a unique_ptr instead of a shared_ptr, for the reason I described above.

But without using PlainVal, we will have to manually do a deep copy of the std::vector<std::unique_ptr<AbstractExprAttribute>> attributes_. It's a little inconvenient, but should be more efficient.

I don't have a strong preference about which one we should choose.

Oh, wait, no. It is not "a little inconvenient", but super convenient. Without knowing T, we can not do a deep copy.

OK, that makes sense. (I assume you meant "super inconvenient")

Let's use PlainVal, but I'd prefer a name like Attribute, which I think a little more clearly implies what it's supposed to be used for.

torch/csrc/jit/codegen/cuda/ir_internal_nodes.h

naoyam · 2022-11-17T23:38:36Z

torch/csrc/jit/codegen/cuda/ir_base_nodes.h

+//! to hold non-IR data, such as DataType, std::vector<int>, etc. Please don't
+//! use this class to hold IR nodes or their pointers.
+template <typename T>
+class TORCH_CUDA_CU_API PlainVal : public Val {


So, this is used to hold IR node attributes. Does it need to be a subclass of Val?

naoyam · 2022-11-17T23:51:36Z

torch/csrc/jit/codegen/cuda/ir_internal_nodes.h

-  const std::vector<Val*>& initVals() const {
-    return init_vals_;
+  std::vector<Val*> initVals() const {
+    return {attributes().begin() + 2, attributes().end()};


A bit concerning that we need to make sure we use correct offsets in the attribute vector. It seems like we are effectively losing static type safety, although the overall IR classes are more consistent.

Given that typically there are not many attributes in each class, I don't think the above concern is critical.

…expr

naoyam

LGTM. Thanks for the refactoring.

zasdfgbnm added 17 commits November 14, 2022 16:51

save structured expr

c0a5864

save ExprType change

ee0f43a

revert

b7aedbd

Merge branch 'devel' of github.com:csarofeen/pytorch into no-expr-type

6149751

more

77be1f4

one of

cd8817c

std::type_index

0ae70ab

more

1f9cb3b

more

cdee6d0

more

8304b84

fix

1e66082

fix reduction

31e61e3

isStrictlyOneOf

7d5b4f7

fix

895b6f0

fix

79017b8

fix

20b5394

Merge branch 'devel' of github.com:csarofeen/pytorch into structured-…

de37d06

…expr

zasdfgbnm changed the base branch from devel to no-expr-type November 15, 2022 23:55

zasdfgbnm added 8 commits November 15, 2022 15:56

Revert "save structured expr"

e4e235a

This reverts commit c0a5864.

Merge branch 'no-expr-type' of github.com:csarofeen/pytorch into stru…

f8fccd0

…ctured-expr

FullOp compiles

31fc54b

ARangeOp

0da63b0

TernaryOpType

9e57950

RNGOp

3fb5d26

BroadcastOp SqueezeOp

73db1ea

ReductionOp

1b2bc68

zasdfgbnm changed the title ~~WIP~~ Reduce the work to add a new expr: unify the structure of exprs Nov 16, 2022

zasdfgbnm changed the title ~~Reduce the work to add a new expr: unify the structure of exprs~~ [WIP] Reduce the work to add a new expr: unify the structure of exprs Nov 16, 2022

zasdfgbnm added 2 commits November 15, 2022 23:42

grouped reduction, welford, grouped welford

f0e4e75

runtime time

d7eb2de

zasdfgbnm added 10 commits November 16, 2022 14:26

gather, view as scalar, view, load store

64bd0d8

DEFINE_CLONE

242228e

Allocate

393fa52

ForLoop

e3f27a2

fix

f72d6b6

GroupedGridReduction

4338807

GridWelford

412420b

all expr done

c7b860e

fix

bfdb742

cleanup

36ada0e

zasdfgbnm commented Nov 17, 2022

View reviewed changes

rewrite IrCloner

56b03ee

zasdfgbnm changed the title ~~[WIP] Reduce the work to add a new expr: unify the structure of exprs~~ Reduce the work to add a new expr: unify the structure of exprs Nov 17, 2022

zasdfgbnm marked this pull request as ready for review November 17, 2022 09:40

zasdfgbnm requested a review from naoyam November 17, 2022 09:44

zasdfgbnm commented Nov 17, 2022

View reviewed changes

naoyam reviewed Nov 17, 2022

View reviewed changes

zasdfgbnm added 6 commits November 17, 2022 21:38

rename

c80587c

Merge branch 'devel' of github.com:csarofeen/pytorch into structured-…

522bfdf

…expr

save

9bd3bc4

fix

0987844

save

c23ea6c

save

0971dfe

naoyam approved these changes Nov 18, 2022

View reviewed changes

Statement attributes_

6fd6504

zasdfgbnm merged commit 50861ec into devel Nov 18, 2022

zasdfgbnm deleted the structured-expr branch November 18, 2022 08:27

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Reduce the work to add a new expr: unify the structure of exprs #2190

Reduce the work to add a new expr: unify the structure of exprs #2190

zasdfgbnm commented Nov 15, 2022 •

edited

Loading

zasdfgbnm Nov 17, 2022

zasdfgbnm Nov 17, 2022 •

edited

Loading

naoyam Nov 17, 2022

zasdfgbnm Nov 17, 2022

naoyam Nov 17, 2022

zasdfgbnm Nov 17, 2022

zasdfgbnm Nov 18, 2022

naoyam Nov 17, 2022

naoyam Nov 17, 2022

zasdfgbnm Nov 17, 2022

naoyam Nov 17, 2022

zasdfgbnm Nov 18, 2022

naoyam Nov 18, 2022

zasdfgbnm Nov 18, 2022

zasdfgbnm Nov 18, 2022

naoyam Nov 18, 2022

zasdfgbnm Nov 18, 2022

naoyam Nov 17, 2022

naoyam Nov 17, 2022

naoyam left a comment

		@@ -15,20 +16,6 @@ class Kernel;

		class IrCloner;

		// Passkey for builder to register properties with statements, and to call

Reduce the work to add a new expr: unify the structure of exprs #2190

Reduce the work to add a new expr: unify the structure of exprs #2190

Conversation

zasdfgbnm commented Nov 15, 2022 • edited Loading

Choose a reason for hiding this comment

zasdfgbnm Nov 17, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

naoyam left a comment

Choose a reason for hiding this comment

zasdfgbnm commented Nov 15, 2022 •

edited

Loading

zasdfgbnm Nov 17, 2022 •

edited

Loading