[Mlir-commits] [mlir] [MLIR][OpenMP] Enable multiple variables for target teams reductions (PR #134903)

Sergio Afonso llvmlistbot at llvm.org
Wed Apr 9 07:47:50 PDT 2025


================
@@ -4688,12 +4688,18 @@ static uint64_t getTypeByteSize(mlir::Type type, const DataLayout &dl) {
 template <typename OpTy>
 static uint64_t getReductionDataSize(OpTy &op) {
   if (op.getNumReductionVars() > 0) {
-    assert(op.getNumReductionVars() == 1 &&
-           "Only 1 reduction variable currently supported");
-    mlir::Type reductionVarTy = op.getReductionVars()[0].getType();
+    SmallVector<omp::DeclareReductionOp> reductions;
+    collectReductionDecls(op, reductions);
+
+    llvm::SmallVector<mlir::Type> members;
+    for (omp::DeclareReductionOp &red : reductions) {
+      members.push_back(red.getType());
+    }
----------------
skatrak wrote:

Nit: Reserve space in advance to potentially avoid multiple allocations and braces.
```suggestion
    llvm::SmallVector<mlir::Type> members;
    members.reserve(reductions.size());
    for (omp::DeclareReductionOp &red : reductions)
      members.push_back(red.getType());
```

https://github.com/llvm/llvm-project/pull/134903


More information about the Mlir-commits mailing list