[CODEGEN] Various bugfixes that make it possible to fuse RNG in a matmul epilogue (#356)

include/triton/codegen/selection/generator.h

@@ -147,12 +147,14 @@ public:
   void visit_store_inst(ir::store_inst*);
   void visit_unmasked_store_inst(ir::unmasked_store_inst*);
   void visit_masked_store_inst(ir::masked_store_inst*);
+  void visit_cat_inst(ir::cat_inst*);
   void visit_reshape_inst(ir::reshape_inst*);
   void visit_splat_inst(ir::splat_inst*);
   void visit_broadcast_inst(ir::broadcast_inst*);
   void visit_downcast_inst(ir::downcast_inst*);
   void visit_exp_inst(ir::exp_inst*);
   void visit_cos_inst(ir::cos_inst*);
+  void visit_umulhi_inst(ir::umulhi_inst* x);
   void visit_sin_inst(ir::sin_inst*);
   void visit_log_inst(ir::log_inst*);
   void visit_get_program_id_inst(ir::get_program_id_inst*);

include/triton/ir/builder.h

@@ -137,6 +137,7 @@ public:
   // Block instruction
   value *create_splat(value *arg, const type::block_shapes_t &shapes);
   value *create_reshape(value *arg, const type::block_shapes_t &shapes);
+  value *create_cat(value *lhs, value *rhs);
   value *create_broadcast(value *arg, const type::block_shapes_t &shapes);
   // Built-in instruction
   value *create_get_program_id(unsigned axis);
@@ -153,6 +154,8 @@ public:
   value *create_reduce(value *A, reduce_inst::op_t op, unsigned axis);
   value *create_select(value *pred, value *if_value, value *else_value);
   // Intrinsics
+  // These have no place in the IR, and hopefully they can be removed at some point
+  value *create_umulhi(value* lhs, value* rhs);
   value *create_copy_to_shared(value *arg);
   value *create_masked_load_async(value *arg, value *mask, value *false_value, load_inst::CACHE_MODIFIER cache);
   value *create_copy_from_shared(value *arg);

include/triton/ir/dispatch.h

@@ -61,6 +61,7 @@ struct dispatch{
 
   // casting ops
   static ir::value *reshape(ir::value *input, shape_t shape, ir::builder *builder);
+  static ir::value *cat(ir::value *lhs, ir::value *rhs, ir::builder *builder);
   static ir::value *broadcast(ir::value *input, shape_t shape, ir::builder *builder);
   static std::tuple<ir::value*, ir::value*> broadcast(ir::value *lhs, ir::value* rhs, ir::builder *builder);
   static ir::value *bitcast(ir::value *input, ir::type *type, ir::builder *builder);
@@ -90,6 +91,7 @@ struct dispatch{
   static ir::value *sum(ir::value *input, unsigned int axis, ir::builder *builder);
 
   // math
+  static ir::value *umulhi(ir::value *x, ir::value *y, ir::builder *builder);
   static ir::value *exp(ir::value *x, ir::builder *builder);
   static ir::value *log(ir::value *x, ir::builder *builder);
   static ir::value *cos(ir::value *x, ir::builder *builder);

include/triton/ir/enums.h

@@ -132,6 +132,7 @@ enum value_id_t: unsigned {
   // retile
   INST_RESHAPE,
   INST_SPLAT,
+  INST_CAT,
   INST_BROADCAST,
   INST_DOWNCAST,
   // builtin
@@ -142,6 +143,7 @@ enum value_id_t: unsigned {
   INST_ATOMIC_EXCH,
   INST_ATOMIC_RMW,
   // math
+  INST_UMULHI,
   INST_EXP,
   INST_COS,
   INST_SIN,

include/triton/ir/instructions.h

@@ -520,6 +520,21 @@ public:
 //                               retile_inst classes
 //===----------------------------------------------------------------------===//
 
+// cat
+
+class cat_inst: public instruction {
+private:
+  std::string repr_impl() const { return "cat"; }
+  cat_inst(value *x, value *y, const std::string &name, instruction *next);
+
+public:
+  static instruction* create(value *lhs, value *rhs,
+                             const std::string &name = "",
+                             instruction *next = nullptr);
+  _TRITON_DEFINE_CLONE(cat_inst)
+  _TRITON_DEFINE_ACCEPT(cat_inst)
+};
+
 // retile
 
 class retile_inst: public unary_inst {
@@ -654,6 +669,17 @@ public:
   static instruction* create(value *ptr, value *cmp, value *val, const std::string &name = "", instruction *next = nullptr);
 };
 
+class umulhi_inst: public builtin_inst {
+private:
+  umulhi_inst(value *lhs, value *rhs, const std::string &name = "", instruction *next = nullptr);
+  std::string repr_impl() const { return "umulhi"; }
+  _TRITON_DEFINE_CLONE(umulhi_inst)
+  _TRITON_DEFINE_ACCEPT(umulhi_inst)
+
+public:
+  static instruction* create(value *lhs, value *rhs, const std::string &name = "", instruction *next = nullptr);
+};
+
 class exp_inst: public builtin_inst {
 private:
   exp_inst(value *val, const std::string &name = "", instruction *next = nullptr);
@@ -803,6 +829,7 @@ public:
 //                               intrinsics classes
 //===----------------------------------------------------------------------===//
 
+
 class copy_to_shared_inst: public unary_inst{
 private:
   using unary_inst::unary_inst;
@@ -884,35 +911,6 @@ public:
    instruction *next=nullptr);
 };
 
-//// On NVIDIA, implementation is such that
-//// constant_range = nv_dynamic_program_idx + nv_static_program_idx
-//// so as to enable re-association on nv_static_program_idx which is constant
-//class make_range_dyn: public instruction {
-//private:
-//  make_range_dyn(type *ty, const std::string &name, instruction *next);
-//  std::string repr_impl() const { return "nv_dynamic_program_idx"; }
-//  _TRITON_DEFINE_CLONE(make_range_dyn)
-//  _TRITON_DEFINE_ACCEPT(make_range_dyn)
-
-//public:
-//  static make_range_dyn* create(type *ty, const std::string &name = "", instruction *next = nullptr);
-//};
-
-//class make_range_sta: public constant {
-//private:
-//  make_range_sta(make_range *range);
-
-//public:
-//  static make_range_sta *get(make_range* range);
-//  make_range* get_range() const;
-//  std::string repr() const { return "nv_static_program_idx"; }
-//  _TRITON_DEFINE_ACCEPT(make_range_sta)
-
-//private:
-//  make_range *range_;
-//};
-
-
 /* constant range */
 class make_range: public instruction{
   make_range(type *ty, constant_int* first, constant_int* last);

include/triton/ir/visitor.h

@@ -45,9 +45,11 @@ class masked_store_inst;
 class retile_inst;
 class reshape_inst;
 class splat_inst;
+class cat_inst;
 class broadcast_inst;
 class downcast_inst;
 
+class umulhi_inst;
 class exp_inst;
 class cos_inst;
 class sin_inst;
@@ -122,6 +124,7 @@ public:
   virtual void visit_unmasked_store_inst(unmasked_store_inst*) = 0;
   virtual void visit_masked_store_inst(masked_store_inst*) = 0;
 
+  virtual void visit_umulhi_inst(umulhi_inst*) = 0;
   virtual void visit_exp_inst(exp_inst*) = 0;
   virtual void visit_cos_inst(cos_inst*) = 0;
   virtual void visit_sin_inst(sin_inst*) = 0;
@@ -129,6 +132,7 @@ public:
 
   virtual void visit_reshape_inst(reshape_inst*) = 0;
   virtual void visit_splat_inst(splat_inst*) = 0;
+  virtual void visit_cat_inst(cat_inst*) = 0;
   virtual void visit_broadcast_inst(broadcast_inst*) = 0;
   virtual void visit_downcast_inst(downcast_inst*) = 0;
 
@@ -150,13 +154,10 @@ public:
   virtual void visit_masked_load_async_inst(masked_load_async_inst*)= 0;
   virtual void visit_barrier_inst(barrier_inst*) = 0;
   virtual void visit_async_wait_inst(async_wait_inst*) = 0;
-//  virtual void visit_make_range_dyn(make_range_dyn*) = 0;
   virtual void visit_make_range(make_range*) = 0;
   virtual void visit_prefetch_s_inst(prefetch_s_inst*) = 0;
-
   virtual void visit_function(function*) = 0;
 
-//  virtual void visit_make_range_sta(make_range_sta*) = 0;
   virtual void visit_undef_value(undef_value*) = 0;
   virtual void visit_constant_int(constant_int*) = 0;
   virtual void visit_constant_fp(constant_fp*) = 0;