From 86d94e02478aaae1e8fc54463db23ed61f2488d8 Mon Sep 17 00:00:00 2001
From: Kavitha Madhu <kavitha.madhu@intel.com>
Date: Tue, 30 Jul 2024 21:58:37 -0700
Subject: [PATCH] In parallel op nesting

---
 lib/TPP/Transforms/LoopExpansion.cpp | 38 ++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/lib/TPP/Transforms/LoopExpansion.cpp b/lib/TPP/Transforms/LoopExpansion.cpp
index 2d4f03fd2..522fe048f 100644
--- a/lib/TPP/Transforms/LoopExpansion.cpp
+++ b/lib/TPP/Transforms/LoopExpansion.cpp
@@ -76,6 +76,20 @@ static LogicalResult loopExpand(scf::ForallOp op, unsigned numOuterParallel) {
     }
   }
 
+  for (auto oper = op.getBody()->getOperations().begin();
+       oper != op.getBody()->getOperations().end(); oper++) {
+    if (isa<scf::InParallelOp>(oper)) {
+      auto nestedOperations = dyn_cast<scf::InParallelOp>(oper);
+      for (auto nestedOper = nestedOperations.begin();
+           nestedOper != nestedOperations.end(); nestedOper++) {
+        if (isa<scf::InParallelOp>(nestedOper)) {
+          LLVM_DEBUG(llvm::dbgs()
+                     << "Serialization of nested parallel ops unsupported");
+          return failure();
+        }
+      }
+    }
+  }
   // Clone instructions at the innermost loop level
   IRMapping mapping;
   for (auto oper = op.getBody()->getOperations().begin();
@@ -99,8 +113,32 @@ static LogicalResult loopExpand(scf::ForallOp op, unsigned numOuterParallel) {
         }
         j++;
       }
+    } else {
+      auto nestedOperations = (dyn_cast<scf::InParallelOp>(oper));
+      for (auto nestedOper = nestedOperations.begin();
+           nestedOper != nestedOperations.end(); nestedOper++) {
+        auto clonedInstr = rewriter.clone(*nestedOper, mapping);
+        nestedOper->replaceAllUsesWith(clonedInstr);
+        int j = 0;
+        for (auto arg : clonedInstr->getOperands()) {
+          for (size_t i = 0; i < op.getInductionVars().size(); i++) {
+            if (arg == op.getInductionVars()[i]) {
+              if (i < numOuterParallel) {
+                clonedInstr->setOperand(
+                    j, parallelOpList[i].getInductionVars()[0]);
+              } else {
+                clonedInstr->setOperand(
+                    j, forOpList[i - numOuterParallel].getInductionVar());
+              }
+              break;
+            }
+          }
+          j++;
+        }
+      }
     }
   }
+
   rewriter.eraseOp(op);
   return success();
 }