[CINN] Composite reduce pass supports horizontal fusion (#72166)

Enigmatisms · web-flow · commit 65befb9a6678 · 2025-04-15T14:13:17.000+08:00
diff --git a/paddle/cinn/optim/realize_composite_reduce_pass.cc b/paddle/cinn/optim/realize_composite_reduce_pass.cc
@@ -55,8 +55,16 @@ struct CompositeTypes : public std::vector<common::Type> {
     this->reserve(2);
   }
 
+  bool operator==(const CompositeTypes& other) const {
+    if (this->type != other.type || other.size() != this->size()) return false;
+    for (size_t i = 0; i < other.size(); i++) {
+      if (this->at(i) != other.at(i)) return false;
+    }
+    return true;
+  }
+
   void Print() const {
-    VLOG(4) << "[CompositeTypes]: ";
+    VLOG(4) << "[CompositeTypes]: " << static_cast<int>(this->type);
     for (auto _t : *this) {
       VLOG(4) << _t;
     }
@@ -113,37 +121,82 @@ void SetInitValue(Store store_stmt,
   }
 }
 
-template <typename StmtType>
-std::set<ir::Buffer> CollectReduceBuffers(const BlockRef& body) {
-  std::set<ir::Buffer> buffers;
-
-  const auto VisitFn = [&](const StmtRef& stmt) {
-    if (!stmt.isa<StmtType>()) return;
-    StmtType store_stmt = stmt.as<StmtType>();
-    if (GetReduceType(store_stmt->value()) != ReduceType::kNone) {
-      buffers.insert(store_stmt->tensor().as_tensor()->buffer);
+/**
+ * This function resolves undefined argidx type, for example:
+ * \code
+ * spatial inner loop (argidx type defined)
+ *    tensor_0[...] = cinn_argmax(tensor_1[...], argidx_f32_i64(tensor_2[...],
+ * index))
+ *
+ * follow up cross thread reduce (argidx type undefined)
+ *    tensor_3[...] = cinn_argmax(tensor_4[...], tensor_5[...])
+ * \endcode
+ * In the above undefined case, we can not extract value type, since both
+ * tensors (4 and 5) in the arguments will be of index type, which lefts the
+ * argidx type undefined. So this function basically checks whether tensor_5 is
+ * in the typed_buffers map. Since cross thread reduction usually follows
+ * spatial inner loop reduction, so normally, tensor_0 and tensor_5 will
+ * normally be the same in one reduce block. and since tensor_0's type is
+ * defined, we can use it to resolve tensor_5 (and thus, the undefined
+ * tensor_3)'s type.
+ */
+std::map<ir::Buffer, CompositeTypes> ResolveUndefinedArgIdxType(
+    std::map<ir::Buffer, CompositeTypes>&& typed_buffers,
+    std::vector<Store>&& stores) {
+  for (const auto& store_stmt : stores) {
+    if (auto call_stmt = store_stmt->value().As<ir::Call>()) {
+      if (call_stmt->name != hlir::pe::kArgmaxFuncName &&
+          call_stmt->name != hlir::pe::kArgminFuncName)
+        continue;
+      auto load_stmt = call_stmt->read_args[1].As<ir::Load>();
+      PADDLE_ENFORCE_NOT_NULL(load_stmt,
+                              ::common::errors::PreconditionNotMet(
+                                  "Non-spatial inner loop arg reduce func call "
+                                  "second argument must be load."));
+      auto it = typed_buffers.find(load_stmt->tensor.as_tensor()->buffer);
+      PADDLE_ENFORCE_NE(it,
+                        typed_buffers.end(),
+                        ::common::errors::PreconditionNotMet(
+                            "Referenced buffer '%s' should be defined.",
+                            load_stmt->tensor.as_tensor()->buffer->name));
+      auto composite_type = it->second;
+      typed_buffers.emplace(store_stmt->tensor().as_tensor()->buffer,
+                            composite_type);
     }
-  };
-
-  ir::stmt::Visit(body, VisitFn, [](auto) {});
-  return buffers;
+  }
+  return typed_buffers;
 }
 
-CompositeTypes DetermineReduceType(const BlockRef& body) {
-  CompositeTypes composite_reduce;
-
-  // the blockref can only have one unique reduce type
-  const auto VisitFn = [&composite_reduce](const StmtRef& stmt) {
-    // to stop recursive visit: (1) rtype is defined
-    // (2) stmt is not a Store stmt
-    if (composite_reduce.type != ReduceType::kNone || !stmt.isa<Store>())
-      return;
+std::map<ir::Buffer, CompositeTypes> CollectTypedReduceBuffers(
+    const BlockRef& body, std::vector<Store>* arg_stores) {
+  std::map<ir::Buffer, CompositeTypes> typed_buffers;
+  const auto VisitFn = [&](const StmtRef& stmt) {
+    if (!stmt.isa<Store>()) return;
     Store store_stmt = stmt.as<Store>();
-    composite_reduce = GetArgReduceUnderlyingType(store_stmt->value());
+    if (GetReduceType(store_stmt->value()) != ReduceType::kNone) {
+      auto it = typed_buffers.find(store_stmt->tensor().as_tensor()->buffer);
+      if (it == typed_buffers.end()) {
+        auto composite_type = GetArgReduceUnderlyingType(store_stmt->value());
+        if (composite_type.type == ReduceType::kNone) {
+          arg_stores->emplace_back(store_stmt);
+        } else {
+          // defined composite type can be immediately stored
+          typed_buffers.emplace(store_stmt->tensor().as_tensor()->buffer,
+                                composite_type);
+        }
+      } else {
+        // check whether we will have conflicted store types
+        PADDLE_ENFORCE_EQ(
+            it->second == GetArgReduceUnderlyingType(store_stmt->value()),
+            true,
+            ::common::errors::PreconditionNotMet(
+                "Composite type conflict detected in the buffer map."));
+      }
+    }
   };
 
   ir::stmt::Visit(body, VisitFn, [](auto) {});
-  return composite_reduce;
+  return typed_buffers;
 }
 
 Store GetStoreOfSchedule(const Schedule& stmt) {
@@ -301,9 +354,10 @@ struct StageReduceResultMutator : public ir::stmt::StmtMutator<> {
 };
 
 struct LoadTypeMutator : public ir::IRMutator<> {
-  explicit LoadTypeMutator(const std::map<ir::Buffer, ir::Type>& buffer2type,
-                           const CompositeTypes& ctype)
-      : buffer2type_(buffer2type), comp_type_(ctype) {}
+  explicit LoadTypeMutator(
+      const std::map<ir::Buffer, std::pair<ir::Type, CompositeTypes>>&
+          buffer2type)
+      : buffer2type_(buffer2type) {}
 
   void operator()(ir::Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
 
@@ -314,10 +368,11 @@ struct LoadTypeMutator : public ir::IRMutator<> {
     auto& buffer = node->tensor.as_tensor()->buffer;
     auto it = buffer2type_.find(buffer);
     if (it != buffer2type_.end()) {
-      ir::Type new_type = GetCompositeReduceType(it->second, comp_type_);
+      const auto& [buffer_type, composite_type] = it->second;
+      ir::Type new_type = GetCompositeReduceType(buffer_type, composite_type);
       node->tensor.as_tensor()->set_type(new_type);
       buffer->dtype = new_type;
-      *expr = ir::Cast::Make(it->second, *expr);
+      *expr = ir::Cast::Make(buffer_type, *expr);
     }
   }
 
@@ -347,25 +402,25 @@ struct LoadTypeMutator : public ir::IRMutator<> {
     }
   }
 
-  const std::map<ir::Buffer, ir::Type>& buffer2type_;
-  const CompositeTypes& comp_type_;
+  const std::map<ir::Buffer, std::pair<ir::Type, CompositeTypes>>& buffer2type_;
 };
 
 void SetBufferType(ir::LoweredFunc func,
-                   const std::set<ir::Buffer>& buffers,
-                   const CompositeTypes& composite_type) {
-  // Make a map from the buffers to their element types, otherwise it's hard to
-  // know a buffer's original type.
-  std::map<ir::Buffer, ir::Type> buffer2type;
-  for (auto& buffer : buffers) {
-    buffer2type.emplace(buffer, buffer->dtype);
+                   const std::map<ir::Buffer, CompositeTypes>& typed_buffers) {
+  // Make a map from the buffers to their element and composite reduce types,
+  // otherwise it's hard to know a buffer's original type. The original type
+  // must be known to perform casting (back) in LoadTypeMutator::Visit()
+  std::map<ir::Buffer, std::pair<ir::Type, CompositeTypes>> buffer2type;
+  for (auto& [buffer, reduce_type] : typed_buffers) {
+    buffer2type.emplace(buffer, std::make_pair(buffer->dtype, reduce_type));
   }
 
   // Set function's temp_bufs type
   for (auto& buffer : func->temp_bufs) {
     auto it = buffer2type.find(buffer);
     if (it != buffer2type.end()) {
-      buffer->dtype = GetCompositeReduceType(it->second, composite_type);
+      const auto& [buffer_type, composite_type] = it->second;
+      buffer->dtype = GetCompositeReduceType(buffer_type, composite_type);
     }
   }
 
@@ -379,7 +434,8 @@ void SetBufferType(ir::LoweredFunc func,
     auto it = buffer2type.find(buffer);
     if (it != buffer2type.end()) {
       ir::Expr new_tensor = ir::ir_utils::IRCopy(store_stmt->tensor());
-      ir::Type new_type = GetCompositeReduceType(it->second, composite_type);
+      const auto& [buffer_type, composite_type] = it->second;
+      ir::Type new_type = GetCompositeReduceType(buffer_type, composite_type);
       new_tensor.as_tensor()->set_type(new_type);
       new_tensor.as_tensor()->buffer->dtype = new_type;
       store_stmt->set_tensor(new_tensor);
@@ -391,7 +447,7 @@ void SetBufferType(ir::LoweredFunc func,
 
     // Set load buffer type
     ir::Expr new_value = ir::ir_utils::IRCopy(store_stmt->value());
-    LoadTypeMutator load_type_mutator(buffer2type, composite_type);
+    LoadTypeMutator load_type_mutator(buffer2type);
     load_type_mutator(&new_value);
     store_stmt->set_value(new_value);
   };
@@ -462,19 +518,21 @@ LogicalResult RealizeCompositeReducePass::Run(ir::LoweredFunc func) {
   StageReduceResultMutator mutator(func);
   mutator(body);
 
-  auto composite_type = DetermineReduceType(body);
-  if (composite_type.type == ReduceType::kNone) {
+  // Step 2. Collect buffers that are used for reduce computation.
+  std::vector<Store> arg_stores;
+  auto typed_buffers = CollectTypedReduceBuffers(body, &arg_stores);
+  if (typed_buffers.empty()) {
+    // not a composite reduce func
     return LogicalResult::success();
   }
-
-  // Step 2. Collect buffers that are used for reduce computation.
-  std::set<ir::Buffer> buffers = CollectReduceBuffers<Store>(body);
+  typed_buffers = ResolveUndefinedArgIdxType(std::move(typed_buffers),
+                                             std::move(arg_stores));
 
   // Step 3. Change the data type of buffers to the corresponding type.
-  SetBufferType(func, buffers, composite_type);
+  SetBufferType(func, typed_buffers);
 
-  // Step 4. Replace the `cinn_reduce_variance` calls to `operator+` in order to
-  //   reuse the cross-thread/block reduction templates.
+  // Step 4. Replace the `cinn_reduce_variance` and `cinn_argmax` calls
+  // in order to reuse the cross-thread/block reduction templates.
   ReplaceReduceExternCall(body);
 
   return LogicalResult::success();
diff --git a/test/ir/pir/cinn/test_cinn_argidx_reduce.py b/test/ir/pir/cinn/test_cinn_argidx_reduce.py
@@ -43,12 +43,14 @@ def func(x):
         x = paddle.randn([128, 256])
         self.eval(func, [x])
 
-    def test_argmax_discrete(self):
-        # discrete reduce
+    def test_argmax_discrete_hfuse(self):
+        # discrete reduce + horizontal fuse
+        bn = paddle.nn.BatchNorm(256)
+
         def func(x):
-            return x[x.argmax(axis=0, keepdim=True), paddle.arange(x.shape[-1])]
+            return bn(x), x.argmax(axis=0)
 
-        x = paddle.randn([66, 256])
+        x = paddle.randn([128, 256])
         self.eval(func, [x])
 
     def test_argmin_all_bc(self):
