[Mlir-commits] [mlir] 9e77a45 - [mlir][OpenMP][NFC] Refactor fillAffinityIteratorLoop (#189418)

[llvmlistbot at llvm.org]

Author: Chi-Chun, Chen
Date: 2026-03-31T11:12:00-05:00
New Revision: 9e77a4593526a04cd2308a724039e94d8fbe9445

URL: https://github.com/llvm/llvm-project/commit/9e77a4593526a04cd2308a724039e94d8fbe9445
DIFF: https://github.com/llvm/llvm-project/commit/9e77a4593526a04cd2308a724039e94d8fbe9445.diff

LOG: [mlir][OpenMP][NFC] Refactor fillAffinityIteratorLoop (#189418)

Extract affinity-specific logic from fillAffinityIteratorLoop into a
callback so that the iterator loop codegen logic can be shared with
other clauses such as depend clause and target clause.

Added: 
    

Modified: 
    mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 281235a0af461..9f8c13bbe9060 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -2473,11 +2473,14 @@ convertIteratorRegion(llvm::Value *linearIV, IteratorInfo &iterInfo,
   return mlir::success();
 }
 
+using IteratorStoreEntryTy =
+    llvm::function_ref<void(llvm::Value *linearIV, mlir::omp::YieldOp yield)>;
+
 static mlir::LogicalResult
-fillAffinityIteratorLoop(mlir::omp::IteratorOp itersOp,
-                         llvm::IRBuilderBase &builder,
-                         mlir::LLVM::ModuleTranslation &moduleTranslation,
-                         llvm::Value *affinityList, IteratorInfo &iterInfo) {
+fillIteratorLoop(mlir::omp::IteratorOp itersOp, llvm::IRBuilderBase &builder,
+                 mlir::LLVM::ModuleTranslation &moduleTranslation,
+                 IteratorInfo &iterInfo, llvm::StringRef loopName,
+                 IteratorStoreEntryTy genStoreEntry) {
   mlir::Region &itersRegion = itersOp.getRegion();
   mlir::Block &iteratorRegionBlock = itersRegion.front();
 
@@ -2494,19 +2497,12 @@ fillAffinityIteratorLoop(mlir::omp::IteratorOp itersOp,
           "failed to convert iterator region", llvm::inconvertibleErrorCode());
     }
 
-    // Extract affinity entry from omp.yield and store into list[linearIV].
     auto yield =
         mlir::dyn_cast<mlir::omp::YieldOp>(iteratorRegionBlock.getTerminator());
     assert(yield && yield.getResults().size() == 1 &&
            "expect omp.yield in iterator region to have one result");
-    auto entryOp =
-        yield.getResults()[0].getDefiningOp<mlir::omp::AffinityEntryOp>();
-    assert(entryOp && "expect yield generate an affinity entry");
 
-    llvm::Value *addr = moduleTranslation.lookupValue(entryOp.getAddr());
-    llvm::Value *len = moduleTranslation.lookupValue(entryOp.getLen());
-    storeAffinityEntry(builder, *moduleTranslation.getOpenMPBuilder(),
-                       affinityList, linearIV, addr, len);
+    genStoreEntry(linearIV, yield);
 
     // Iterator-region block/value mappings are temporary for this conversion,
     // clear them to avoid stale entries in ModuleTranslation.
@@ -2517,8 +2513,7 @@ fillAffinityIteratorLoop(mlir::omp::IteratorOp itersOp,
 
   llvm::OpenMPIRBuilder::InsertPointOrErrorTy afterIP =
       moduleTranslation.getOpenMPBuilder()->createIteratorLoop(
-          loc, iterInfo.getTotalTrips(), bodyGen,
-          /*Name=*/"iterator");
+          loc, iterInfo.getTotalTrips(), bodyGen, loopName);
   if (failed(handleError(afterIP, *itersOp)))
     return failure();
 
@@ -2575,8 +2570,20 @@ buildAffinityData(mlir::omp::TaskOp &taskOp, llvm::IRBuilderBase &builder,
       assert(itersOp && "iterated value must be defined by omp.iterator");
       IteratorInfo iterInfo(itersOp, moduleTranslation, builder);
       llvm::Value *affList = allocateAffinityList(iterInfo.getTotalTrips());
-      if (failed(fillAffinityIteratorLoop(itersOp, builder, moduleTranslation,
-                                          affList, iterInfo)))
+      if (failed(fillIteratorLoop(
+              itersOp, builder, moduleTranslation, iterInfo, "iterator",
+              [&](llvm::Value *linearIV, mlir::omp::YieldOp yield) {
+                auto entryOp = yield.getResults()[0]
+                                   .getDefiningOp<mlir::omp::AffinityEntryOp>();
+                assert(entryOp && "expect yield produce an affinity entry");
+                llvm::Value *addr =
+                    moduleTranslation.lookupValue(entryOp.getAddr());
+                llvm::Value *len =
+                    moduleTranslation.lookupValue(entryOp.getLen());
+                storeAffinityEntry(builder,
+                                   *moduleTranslation.getOpenMPBuilder(),
+                                   affList, linearIV, addr, len);
+              })))
         return llvm::failure();
       ads.emplace_back(createAffinity(iterInfo.getTotalTrips(), affList));
     }