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

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

@@ -2003,10 +2003,14 @@ class CudaKernelGenerator : private OptOutConstDispatch {
     ArgumentBuilder read_preds;
     ArgumentBuilder write_preds;
 
+    auto output_vals = grouped_gwop->outputVals();
+    auto input_vals = grouped_gwop->inputVals();
+    auto init_vals = grouped_gwop->initVals();
+
     for (const auto expr_index : c10::irange(grouped_gwop->numExprs())) {
-      const auto& output = grouped_gwop->outputVals().at(expr_index);
-      const auto& input = grouped_gwop->inputVals().at(expr_index);
-      const auto& init = grouped_gwop->initVals().at(expr_index);
+      const auto& output = output_vals.at(expr_index);
+      const auto& input = input_vals.at(expr_index);
+      const auto& init = init_vals.at(expr_index);
 
       for (const auto& group_index :
            c10::irange(index_replacement_maps.size())) {

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

@@ -313,6 +313,10 @@ void Val::constDispatch(T handler, const Val* val) {
     case ValType::TensorIndex:
       ptr(handler)->handle(val->as<kir::TensorIndex>());
       return;
+    case ValType::Attribute:
+      // Attribute Val is just a wrapper for non-IR data, so there is nothing to
+      // handle
+      return;
     default:
       break;
   }

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

@@ -34,6 +34,8 @@ Statement::Statement(const Statement* src, IrCloner* ir_cloner) {
   ir_container_ = ir_cloner->container();
 }
 
+NVFUSER_DEFINE_CLONE(Statement)
+
 void Statement::setName(IrContainerPasskey, StmtNameType name) {
   name_ = name;
 }
@@ -98,6 +100,8 @@ Val::Val(IrBuilderPasskey passkey, ValType _vtype, DataType _dtype)
 Val::Val(const Val* src, IrCloner* ir_cloner)
     : Statement(src, ir_cloner), vtype_(src->vtype_), dtype_(src->dtype_) {}
 
+NVFUSER_DEFINE_CLONE(Val)
+
 const std::vector<Expr*>& Val::uses() const {
   if (vtype_ == ValType::TensorView) {
     if (!fusion()->isTVUseInfoValid() && !fusion()->isUpdatingTVUseInfo()) {
@@ -319,7 +323,27 @@ Expr::Expr(IrBuilderPasskey passkey) : Statement(passkey) {}
 Expr::Expr(const Expr* src, IrCloner* ir_cloner)
     : Statement(src, ir_cloner),
       inputs_(ir_cloner->clone(src->inputs_)),
-      outputs_(ir_cloner->clone(src->outputs_)) {}
+      outputs_(ir_cloner->clone(src->outputs_)),
+      attributes_(ir_cloner->clone(src->attributes_)) {}
+
+Expr::Expr(
+    IrBuilderPasskey passkey,
+    std::vector<Val*> inputs,
+    std::vector<Val*> outputs,
+    std::vector<Statement*> attributes)
+    : Statement(passkey),
+      inputs_(std::move(inputs)),
+      outputs_(std::move(outputs)),
+      attributes_(std::move(attributes)) {}
+
+Expr* Expr::shallowCopy() const {
+  auto result = newObject(inputs(), outputs(), attributes());
+  if (container()->isA<kir::Kernel>()) {
+    result->predicate_ = predicate_;
+    result->write_predicate_ = write_predicate_;
+  }
+  return result;
+}
 
 bool Expr::sameAs(const Statement* other) const {
   if (this == other) {
@@ -333,14 +357,20 @@ bool Expr::sameAs(const Statement* other) const {
     return false;
   }
   if (inputs().size() != other_expr->inputs().size() ||
-      outputs().size() != other_expr->outputs().size()) {
+      outputs().size() != other_expr->outputs().size() ||
+      attributes().size() != other_expr->attributes().size()) {
     return false;
   }
   for (const auto i : c10::irange(inputs().size())) {
     if (!input(i)->sameAs(other_expr->input(i))) {
       return false;
     }
   }
+  for (const auto i : c10::irange(attributes().size())) {
+    if (!attribute(i)->sameAs(other_expr->attribute(i))) {
+      return false;
+    }
+  }
   return true;
 }
 
@@ -380,13 +410,6 @@ Expr* Expr::withWritePredicate(kir::Predicate* predicate) {
   return result;
 }
 
-void Expr::copyPredicatesFrom(const Expr* expr) {
-  if (container()->isA<kir::Kernel>()) {
-    predicate_ = expr->predicate_;
-    write_predicate_ = expr->write_predicate_;
-  }
-}
-
 } // namespace cuda
 } // namespace fuser
 } // namespace jit

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

@@ -5,6 +5,7 @@
 #include <c10/util/Exception.h>
 #include <c10/util/Optional.h>
 
+#include <torch/csrc/jit/codegen/cuda/ir_builder_passkey.h>
 #include <torch/csrc/jit/codegen/cuda/type.h>
 #include <torch/csrc/jit/codegen/cuda/utils.h>
 
@@ -70,6 +71,14 @@ class ExprPasskey {
 
 TORCH_CUDA_CU_API void swap(Fusion& a, Fusion& b) noexcept;
 
+#define NVFUSER_DECLARE_CLONE \
+  virtual Statement* clone(IrCloner* ir_cloner) const override;
+
+#define NVFUSER_DEFINE_CLONE(ClassName)                    \
+  Statement* ClassName::clone(IrCloner* ir_cloner) const { \
+    return IrBuilder::clone(this, ir_cloner);              \
+  }
+
 //! Statement is the highest level node representation. Everything that is
 //! considered "IR" will be derived from this class at some point. Both Values
 //! and Expr's are a Statement. If there will ever be any more fundamental
@@ -159,6 +168,8 @@ class TORCH_CUDA_CU_API Statement : public NonCopyable, public PolymorphicBase {
   std::string toString() const;
   std::string toInlineString() const;
 
+  virtual Statement* clone(IrCloner* ir_cloner) const;
+
  protected:
   Statement(IrBuilderPasskey);
 
@@ -353,6 +364,8 @@ class TORCH_CUDA_CU_API Val : public Statement {
 
   void resolveIndexDtype();
 
+  NVFUSER_DECLARE_CLONE
+
  protected:
   friend Fusion;
 
@@ -391,6 +404,31 @@ class TORCH_CUDA_CU_API Val : public Statement {
   int evaluator_index_ = -1;
 };
 
+//! A Val object that stores a plain data. Note that this class is only intended
+//! 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 Attribute : public Val {
+ public:
+  T value;
+  Attribute(IrBuilderPasskey passkey, const T& value)
+      : Val(passkey, ValType::Attribute), value(value) {}
+  Attribute(const Attribute* src, IrCloner* ir_cloner)
+      : Val(src, ir_cloner), value(src->value) {}
+  template <typename... Args>
+  Attribute(IrBuilderPasskey passkey, Args... args)
+      : Val(passkey, ValType::Attribute), value(std::forward<Args>(args)...) {}
+
+  NVFUSER_DECLARE_CLONE
+
+  bool sameAs(const Statement* other) const override {
+    if (auto pv = dynamic_cast<const Attribute*>(other)) {
+      return pv->value == value;
+    }
+    return false;
+  }
+};
+
 //!  A Expr represents a "computation." These are functions that takes inputs
 //!  and produce outputs, inputs and outputs all being Vals. There are
 //!  specializations of BinaryOp which takes 2 inputs and produces 1 output, and
@@ -436,12 +474,25 @@ class TORCH_CUDA_CU_API Expr : public Statement {
 
   Expr(const Expr* src, IrCloner* ir_cloner);
 
+  Expr(
+      IrBuilderPasskey,
+      std::vector<Val*> inputs,
+      std::vector<Val*> outputs,
+      std::vector<Statement*> attributes);
+
   // 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;
 
   bool sameAs(const Statement* other) const override;
 
+  // Creates a new instance of the same expression type with the given inputs,
+  // outputs, and attributes.
+  virtual Expr* newObject(
+      std::vector<Val*> inputs,
+      std::vector<Val*> outputs,
+      std::vector<Statement*> attributes) const = 0;
+
   // Input/output accessors
   const auto& inputs() const {
     return inputs_;
@@ -451,12 +502,24 @@ class TORCH_CUDA_CU_API Expr : public Statement {
     return outputs_;
   }
 
+  const auto& attributes() const {
+    return attributes_;
+  }
+
   auto input(size_t index) const {
-    return inputs_[index];
+    return inputs_.at(index);
   }
 
   auto output(size_t index) const {
-    return outputs_[index];
+    return outputs_.at(index);
+  }
+
+  auto attribute(size_t index) const {
+    return attributes_.at(index);
+  }
+
+  auto attributeVal(size_t index) const {
+    return dynamic_cast<Val*>(attributes_.at(index));
   }
 
   // Dispatch functions, definitions in dispatch.cpp
@@ -494,8 +557,6 @@ class TORCH_CUDA_CU_API Expr : public Statement {
   // TODO: Protect based on being in kernel container
   void setWritePredicate(kir::Predicate* write_predicate);
 
-  void copyPredicatesFrom(const Expr* expr);
-
   // TODO: Add Fusion passkey
   void addInput(Val* input) {
     TORCH_INTERNAL_ASSERT(input != nullptr);
@@ -508,13 +569,19 @@ class TORCH_CUDA_CU_API Expr : public Statement {
     outputs_.push_back(output);
   }
 
+  // TODO: Add Fusion passkey
+  void addAttribute(Statement* attr) {
+    attributes_.push_back(attr);
+  }
+
   ExprPasskey exprPasskey() {
     return ExprPasskey();
   }
 
  private:
   std::vector<Val*> inputs_;
   std::vector<Val*> outputs_;
+  std::vector<Statement*> attributes_;
 
   kir::Predicate* predicate_ = nullptr;
 
@@ -530,6 +597,24 @@ bool Val::isDefinitionType() const {
   return false;
 }
 
+#define NVFUSER_DECLARE_CLONE_AND_CREATE                        \
+  virtual Statement* clone(IrCloner* ir_cloner) const override; \
+  virtual Expr* newObject(                                      \
+      std::vector<Val*> inputs,                                 \
+      std::vector<Val*> outputs,                                \
+      std::vector<Statement*> attributes) const override;
+
+#define NVFUSER_DEFINE_CLONE_AND_CREATE(ClassName)                    \
+  Statement* ClassName::clone(IrCloner* ir_cloner) const {            \
+    return IrBuilder::clone(this, ir_cloner);                         \
+  }                                                                   \
+  Expr* ClassName::newObject(                                         \
+      std::vector<Val*> inputs,                                       \
+      std::vector<Val*> outputs,                                      \
+      std::vector<Statement*> attributes) const {                     \
+    return IrBuilder::create<ClassName>(inputs, outputs, attributes); \
+  }
+
 } // namespace cuda
 } // namespace fuser
 } // namespace jit

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

@@ -11,75 +11,6 @@ namespace jit {
 namespace fuser {
 namespace cuda {
 
-//! Clone an IR node, forwarding the arguments to the IrCloner constructor.
-template <class T>
-T* IrBuilder::clone(const T* src, IrCloner* ir_cloner) {
-  TORCH_INTERNAL_ASSERT(
-      ir_cloner != nullptr,
-      "Cannot use create when a cloner object is set. Use clone.");
-
-  TORCH_INTERNAL_ASSERT(
-      ir_cloner->container() != nullptr,
-      "Cloner doesn't have a valid container to store cloned object.");
-
-  T* dest = new T(src, ir_cloner);
-  const Statement* src_stmt = dynamic_cast<const Statement*>(src);
-  Statement* dest_stmt = dynamic_cast<Statement*>(dest);
-
-  auto dest_container = ir_cloner->container();
-  auto src_container = src_stmt->container();
-
-  dest_container->registerStmt(IrBuilderPasskey(dest_container), dest_stmt);
-
-  if (src_container != dest_container) {
-    dest_stmt->setName(IrBuilderPasskey(dest_container), src_stmt->name());
-  }
-
-  ir_cloner->registerClone(src_stmt, dest_stmt);
-
-  return dest;
-}
-
-#define IR_BUILDER_INSTANTIATE(T) \
-  template T* IrBuilder::clone(const T* src, IrCloner* ir_cloner);
-
-// Vals
-IR_BUILDER_INSTANTIATE(IterDomain)
-IR_BUILDER_INSTANTIATE(TensorDomain)
-IR_BUILDER_INSTANTIATE(TensorView)
-IR_BUILDER_INSTANTIATE(Bool)
-IR_BUILDER_INSTANTIATE(Float)
-IR_BUILDER_INSTANTIATE(Double)
-IR_BUILDER_INSTANTIATE(Int)
-IR_BUILDER_INSTANTIATE(ComplexDouble)
-IR_BUILDER_INSTANTIATE(NamedScalar)
-
-// Exprs
-IR_BUILDER_INSTANTIATE(Split)
-IR_BUILDER_INSTANTIATE(Merge)
-IR_BUILDER_INSTANTIATE(Swizzle2D)
-IR_BUILDER_INSTANTIATE(TransposeOp)
-IR_BUILDER_INSTANTIATE(ExpandOp)
-IR_BUILDER_INSTANTIATE(ShiftOp)
-IR_BUILDER_INSTANTIATE(GatherOp)
-IR_BUILDER_INSTANTIATE(ViewAsScalar)
-IR_BUILDER_INSTANTIATE(ViewOp)
-IR_BUILDER_INSTANTIATE(FullOp)
-IR_BUILDER_INSTANTIATE(ARangeOp)
-IR_BUILDER_INSTANTIATE(EyeOp)
-IR_BUILDER_INSTANTIATE(UnaryOp)
-IR_BUILDER_INSTANTIATE(BinaryOp)
-IR_BUILDER_INSTANTIATE(TernaryOp)
-IR_BUILDER_INSTANTIATE(SelectOp)
-IR_BUILDER_INSTANTIATE(RNGOp)
-IR_BUILDER_INSTANTIATE(ReductionOp)
-IR_BUILDER_INSTANTIATE(GroupedReductionOp)
-IR_BUILDER_INSTANTIATE(WelfordOp)
-IR_BUILDER_INSTANTIATE(LoadStoreOp)
-IR_BUILDER_INSTANTIATE(MmaOp)
-IR_BUILDER_INSTANTIATE(BroadcastOp)
-IR_BUILDER_INSTANTIATE(SqueezeOp)
-
 Val* IrBuilder::newResult(DataType dtype) {
   switch (dtype) {
     case DataType::Bool:

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

@@ -111,6 +111,9 @@ class TORCH_CUDA_CU_API SimplifyingIrBuilder : public IrBuilder {
   static Val* minExpr(Val* lhs, Val* rhs);
 };
 
+template <DataType DT>
+NVFUSER_DEFINE_CLONE(FloatingPoint<DT>)
+
 } // namespace cuda
 } // namespace fuser
 } // namespace jit