[MLIR][OpenMP] Created callback function for adding the metadata of nontemporal and handled the translation in OpenMPToLLVMIRTranslation.cpp

Author: kaviya2510

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

@@ -1229,6 +1229,9 @@ class OpenMPIRBuilder {
   void unrollLoopPartial(DebugLoc DL, CanonicalLoopInfo *Loop, int32_t Factor,
                          CanonicalLoopInfo **UnrolledCLI);
 
+  using NonTemporalBodyGenCallbackTy =
+      function_ref<void(llvm::BasicBlock *BB, MDNode *NontemporalNode)>;
+
   /// Add metadata to simd-ize a loop. If IfCond is not nullptr, the loop
   /// is cloned. The metadata which prevents vectorization is added to
   /// to the cloned loop. The cloned loop is executed when ifCond is evaluated
@@ -1242,10 +1245,15 @@ class OpenMPIRBuilder {
   /// \param Order       The enum to map order clause.
   /// \param Simdlen     The Simdlen length to apply to the simd loop.
   /// \param Safelen     The Safelen length to apply to the simd loop.
-  void applySimd(CanonicalLoopInfo *Loop,
-                 MapVector<Value *, Value *> AlignedVars, Value *IfCond,
-                 omp::OrderKind Order, ConstantInt *Simdlen,
-                 ConstantInt *Safelen, ArrayRef<Value *> NontempralVars = {});
+  /// \param NontemporalCBFunc    Call back function for nontemporal.
+  /// \param NontemporalVars      Array of nontemporal vars.
+  void applySimd(
+      CanonicalLoopInfo *Loop, MapVector<Value *, Value *> AlignedVars,
+      Value *IfCond, omp::OrderKind Order, ConstantInt *Simdlen,
+      ConstantInt *Safelen,
+      NonTemporalBodyGenCallbackTy NontemporalCBFunc = [](BasicBlock *,
+                                                          MDNode *) {},
+      ArrayRef<Value *> NontempralVars = {});
 
   /// Generator for '#omp flush'
   ///

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -5385,86 +5385,11 @@ OpenMPIRBuilder::getOpenMPDefaultSimdAlign(const Triple &TargetTriple,
   return 0;
 }
 
-static void appendNontemporalVars(BasicBlock *Block,
-                                  SmallVectorImpl<Value *> &NontemporalVars) {
-  for (Instruction &I : *Block) {
-    if (const CallInst *CI = dyn_cast<CallInst>(&I)) {
-      if (CI->getIntrinsicID() == Intrinsic::memcpy) {
-        llvm::Value *DestPtr = CI->getArgOperand(0);
-        llvm::Value *SrcPtr = CI->getArgOperand(1);
-        for (const llvm::Value *Var : NontemporalVars) {
-          if (Var == SrcPtr) {
-            NontemporalVars.push_back(DestPtr);
-            break;
-          }
-        }
-      }
-    }
-  }
-}
-
-/** Attach nontemporal metadata to the load/store instructions of nontemporal
- *  variables of \p Block
- * Nontemporal variables may be a scalar, fixed size or allocatable
- * or pointer array
- *
- * Example scenarios for nontemporal variables:
- * Case 1: Scalar variable
- *    If the nontemporal variable is a scalar, it is allocated on stack.Load and
- * store instructions directly access the alloca pointer of the scalar
- * variable for fetching information about scalar variable or  writing
- * into the scalar variable. Mark those load and store instructions as
- * non-temporal.
- *
- * Case 2: Fixed Size array
- *    If the nontemporal variable is a fixed-size array, it is allocated
- * as a contiguous block of memory. It uses one GEP instruction, to compute the
- * address of each individual array elements and perform load or store
- * operation on it. Mark those load and store instructions as non-temporal.
- *
- * Case 3: Allocatable array
- *     For an allocatable array, which might involve runtime type descriptor,
- * needs to navigate through descriptors using two or more GEP and load
- * instructions to compute the address of each individual element in an array.
- * Mark those load or store which access the individual array elements as
- * non-temporal.
- */
-static void addNonTemporalMetadata(BasicBlock *Block, MDNode *Nontemporal,
-                                   SmallVectorImpl<Value *> &NontemporalVars) {
-  appendNontemporalVars(Block, NontemporalVars);
-  for (Instruction &I : *Block) {
-    llvm::Value *mem_ptr = nullptr;
-    bool MetadataFlag = true;
-    if (llvm::LoadInst *li = dyn_cast<llvm::LoadInst>(&I)) {
-      if (!(li->getType()->isPointerTy()))
-        mem_ptr = li->getPointerOperand();
-    } else if (llvm::StoreInst *si = dyn_cast<llvm::StoreInst>(&I))
-      mem_ptr = si->getPointerOperand();
-    if (mem_ptr) {
-      while (mem_ptr && !(isa<llvm::AllocaInst>(mem_ptr))) {
-        if (llvm::GetElementPtrInst *gep =
-                dyn_cast<llvm::GetElementPtrInst>(mem_ptr)) {
-          llvm::Type *sourceType = gep->getSourceElementType();
-          if (sourceType->isStructTy() && gep->getNumIndices() >= 2 &&
-              !(gep->hasAllZeroIndices())) {
-            MetadataFlag = false;
-            break;
-          }
-          mem_ptr = gep->getPointerOperand();
-        } else if (llvm::LoadInst *li = dyn_cast<llvm::LoadInst>(mem_ptr))
-          mem_ptr = li->getPointerOperand();
-      }
-      if (MetadataFlag && is_contained(NontemporalVars, mem_ptr))
-        I.setMetadata(LLVMContext::MD_nontemporal, Nontemporal);
-    }
-  }
-}
-
-void OpenMPIRBuilder::applySimd(CanonicalLoopInfo *CanonicalLoop,
-                                MapVector<Value *, Value *> AlignedVars,
-                                Value *IfCond, OrderKind Order,
-                                ConstantInt *Simdlen, ConstantInt *Safelen,
-                                ArrayRef<Value *> NontemporalVarsIn) {
+void OpenMPIRBuilder::applySimd(
+    CanonicalLoopInfo *CanonicalLoop, MapVector<Value *, Value *> AlignedVars,
+    Value *IfCond, OrderKind Order, ConstantInt *Simdlen, ConstantInt *Safelen,
+    OpenMPIRBuilder::NonTemporalBodyGenCallbackTy NontemporalCBFunc,
+    ArrayRef<Value *> NontemporalVarsIn) {
   LLVMContext &Ctx = Builder.getContext();
 
   Function *F = CanonicalLoop->getFunction();
@@ -5562,12 +5487,12 @@ void OpenMPIRBuilder::applySimd(CanonicalLoopInfo *CanonicalLoop,
   }
 
   addLoopMetadata(CanonicalLoop, LoopMDList);
-  SmallVector<Value *> NontemporalVars{NontemporalVarsIn};
+
   // Set nontemporal metadata to load and stores of nontemporal values
-  if (NontemporalVars.size()) {
+  if (NontemporalVarsIn.size()) {
     MDNode *NontemporalNode = MDNode::getDistinct(Ctx, {});
     for (BasicBlock *BB : Reachable)
-      addNonTemporalMetadata(BB, NontemporalNode, NontemporalVars);
+      NontemporalCBFunc(BB, NontemporalNode);
   }
 }
 

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

@@ -2462,6 +2462,25 @@ convertOrderKind(std::optional<omp::ClauseOrderKind> o) {
   llvm_unreachable("Unknown ClauseOrderKind kind");
 }
 
+static void
+appendNontemporalVars(llvm::BasicBlock *Block,
+                      SmallVectorImpl<llvm::Value *> &NontemporalVars) {
+  for (llvm::Instruction &I : *Block) {
+    if (const llvm::CallInst *CI = dyn_cast<llvm::CallInst>(&I)) {
+      if (CI->getIntrinsicID() == llvm::Intrinsic::memcpy) {
+        llvm::Value *DestPtr = CI->getArgOperand(0);
+        llvm::Value *SrcPtr = CI->getArgOperand(1);
+        for (const llvm::Value *Var : NontemporalVars) {
+          if (Var == SrcPtr) {
+            NontemporalVars.push_back(DestPtr);
+            break;
+          }
+        }
+      }
+    }
+  }
+}
+
 /// Converts an OpenMP simd loop into LLVM IR using OpenMPIRBuilder.
 static LogicalResult
 convertOmpSimd(Operation &opInst, llvm::IRBuilderBase &builder,
@@ -2523,13 +2542,71 @@ convertOmpSimd(Operation &opInst, llvm::IRBuilderBase &builder,
   llvm::MapVector<llvm::Value *, llvm::Value *> alignedVars;
   llvm::omp::OrderKind order = convertOrderKind(simdOp.getOrder());
 
-  llvm::SmallVector<llvm::Value *> nontemporalVars;
+  llvm::SmallVector<llvm::Value *> nontemporalOrigVars;
   mlir::OperandRange nontemporals = simdOp.getNontemporalVars();
   for (mlir::Value nontemporal : nontemporals) {
     llvm::Value *nt = moduleTranslation.lookupValue(nontemporal);
-    nontemporalVars.push_back(nt);
+    nontemporalOrigVars.push_back(nt);
   }
 
+  /** Call back function to attach nontemporal metadata to the load/store
+   * instructions of nontemporal variables of Block.
+   * Nontemporal variables may be a scalar, fixed size or allocatable
+   * or pointer array
+   *
+   * Example scenarios for nontemporal variables:
+   * Case 1: Scalar variable
+   *    If the nontemporal variable is a scalar, it is allocated on stack.Load
+   * and store instructions directly access the alloca pointer of the scalar
+   * variable for fetching information about scalar variable or  writing
+   * into the scalar variable. Mark those load and store instructions as
+   * non-temporal.
+   *
+   * Case 2: Fixed Size array
+   *    If the nontemporal variable is a fixed-size array, it is allocated
+   * as a contiguous block of memory. It uses one GEP instruction, to compute
+   * the address of each individual array elements and perform load or store
+   * operation on it. Mark those load and store instructions as non-temporal.
+   *
+   * Case 3: Allocatable array
+   *     For an allocatable array, which might involve runtime type descriptor,
+   * needs to navigate through descriptors using two or more GEP and load
+   * instructions to compute the address of each individual element in an array.
+   * Mark those load or store which access the individual array elements as
+   * non-temporal.
+   */
+  auto addNonTemporalMetadataCB = [&](llvm::BasicBlock *Block,
+                                      llvm::MDNode *Nontemporal) {
+    SmallVector<llvm::Value *> NontemporalVars{nontemporalOrigVars};
+    appendNontemporalVars(Block, NontemporalVars);
+    for (llvm::Instruction &I : *Block) {
+      llvm::Value *mem_ptr = nullptr;
+      bool MetadataFlag = true;
+      if (llvm::LoadInst *li = dyn_cast<llvm::LoadInst>(&I)) {
+        if (!(li->getType()->isPointerTy()))
+          mem_ptr = li->getPointerOperand();
+      } else if (llvm::StoreInst *si = dyn_cast<llvm::StoreInst>(&I))
+        mem_ptr = si->getPointerOperand();
+      if (mem_ptr) {
+        while (mem_ptr && !(isa<llvm::AllocaInst>(mem_ptr))) {
+          if (llvm::GetElementPtrInst *gep =
+                  dyn_cast<llvm::GetElementPtrInst>(mem_ptr)) {
+            llvm::Type *sourceType = gep->getSourceElementType();
+            if (sourceType->isStructTy() && gep->getNumIndices() >= 2 &&
+                !(gep->hasAllZeroIndices())) {
+              MetadataFlag = false;
+              break;
+            }
+            mem_ptr = gep->getPointerOperand();
+          } else if (llvm::LoadInst *li = dyn_cast<llvm::LoadInst>(mem_ptr))
+            mem_ptr = li->getPointerOperand();
+        }
+        if (MetadataFlag && is_contained(NontemporalVars, mem_ptr))
+          I.setMetadata(llvm::LLVMContext::MD_nontemporal, Nontemporal);
+      }
+    }
+  };
+
   llvm::BasicBlock *sourceBlock = builder.GetInsertBlock();
   std::optional<ArrayAttr> alignmentValues = simdOp.getAlignments();
   mlir::OperandRange operands = simdOp.getAlignedVars();
@@ -2557,11 +2634,11 @@ convertOmpSimd(Operation &opInst, llvm::IRBuilderBase &builder,
 
   builder.SetInsertPoint(*regionBlock, (*regionBlock)->begin());
   llvm::CanonicalLoopInfo *loopInfo = findCurrentLoopInfo(moduleTranslation);
-  ompBuilder->applySimd(loopInfo, alignedVars,
-                        simdOp.getIfExpr()
-                            ? moduleTranslation.lookupValue(simdOp.getIfExpr())
-                            : nullptr,
-                        order, simdlen, safelen, nontemporalVars);
+  ompBuilder->applySimd(
+      loopInfo, alignedVars,
+      simdOp.getIfExpr() ? moduleTranslation.lookupValue(simdOp.getIfExpr())
+                         : nullptr,
+      order, simdlen, safelen, addNonTemporalMetadataCB, nontemporalOrigVars);
 
   return cleanupPrivateVars(builder, moduleTranslation, simdOp.getLoc(),
                             llvmPrivateVars, privateDecls);