Add getIntrinsicIDAndArgs function

Author: clementval

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

@@ -1027,21 +1027,22 @@ def NVVM_BarrierOp : NVVM_Op<"barrier", [AttrSizedOperandSegments]> {
     [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-bar)
   }];
 
+  let extraClassDeclaration = [{
+    static mlir::NVVM::IDArgPair
+      getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
+                            llvm::IRBuilderBase& builder);
+  }];
+
   let arguments = (ins Optional<I32>:$barrierId, Optional<I32>:$numberOfThreads,
       OptionalAttr<BarrierReductionAttr>:$reductionOp,
       Optional<I32>:$reductionOperand);
   string llvmBuilder = [{
-    llvm::Value *id = $barrierId ? $barrierId : builder.getInt32(0);
-    if ($numberOfThreads)
-      createIntrinsicCall(
-          builder, llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_count,
-          {id, $numberOfThreads});
-    else if ($reductionOp)
-      createIntrinsicCall(
-          builder, getBarrierIntrinsicID($reductionOp), {$reductionOperand});
-    else
-      createIntrinsicCall(
-          builder, llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_all, {id});
+    auto [id, args] = NVVM::BarrierOp::getIntrinsicIDAndArgs(
+                        *op, moduleTranslation, builder);
+    if ($reductionOp)
+      $res = createIntrinsicCall(builder, id, args);
+    else 
+      createIntrinsicCall(builder, id, args);
   }];
   let results = (outs Optional<I32>:$res);
 

mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp

@@ -1505,11 +1505,11 @@ LogicalResult NVVM::BarrierOp::verify() {
     return emitOpError(
         "barrier id is missing, it should be set between 0 to 15");
 
-  if (getBarrierId() && (getReductionOp() || getReductionPredicate()))
+  if (getBarrierId() && (getReductionOp() || getReductionOperand()))
     return emitOpError("reduction are only available when id is 0");
 
-  if ((getReductionOp() && !getReductionPredicate()) ||
-      (!getReductionOp() && getReductionPredicate()))
+  if ((getReductionOp() && !getReductionOperand()) ||
+      (!getReductionOp() && getReductionOperand()))
     return emitOpError("reduction predicate and reduction operation must be "
                        "specified together");
 
@@ -1770,6 +1770,41 @@ static bool isPtrInSharedCTASpace(mlir::Value ptr) {
   return isPtrInAddrSpace(ptr, NVVMMemorySpace::Shared);
 }
 
+mlir::NVVM::IDArgPair NVVM::BarrierOp::getIntrinsicIDAndArgs(
+    Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
+  auto thisOp = cast<NVVM::BarrierOp>(op);
+  llvm::Value *barrierId = thisOp.getBarrierId()
+                               ? mt.lookupValue(thisOp.getBarrierId())
+                               : builder.getInt32(0);
+  llvm::Intrinsic::ID id;
+  llvm::SmallVector<llvm::Value *> args;
+  if (thisOp.getNumberOfThreads()) {
+    id = llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_count;
+    args.push_back(barrierId);
+    args.push_back(mt.lookupValue(thisOp.getNumberOfThreads()));
+  } else if (thisOp.getReductionOp()) {
+    switch (*thisOp.getReductionOp()) {
+    case NVVM::BarrierReduction::AND:
+      id = llvm::Intrinsic::nvvm_barrier0_and;
+      break;
+    case NVVM::BarrierReduction::OR:
+      id = llvm::Intrinsic::nvvm_barrier0_or;
+      break;
+    case NVVM::BarrierReduction::POPC:
+      id = llvm::Intrinsic::nvvm_barrier0_popc;
+      break;
+    default:
+      llvm_unreachable("Unknown reduction operation for barrier");
+    }
+    args.push_back(mt.lookupValue(thisOp.getReductionOperand()));
+  } else {
+    id = llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_all;
+    args.push_back(barrierId);
+  }
+
+  return {id, std::move(args)};
+}
+
 mlir::NVVM::IDArgPair MBarrierInitOp::getIntrinsicIDAndArgs(
     Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
   auto thisOp = cast<NVVM::MBarrierInitOp>(op);

mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp

@@ -291,22 +291,6 @@ static unsigned getUnidirectionalFenceProxyID(NVVM::ProxyKind fromProxy,
   llvm_unreachable("Unsupported proxy kinds");
 }
 
-static unsigned
-getBarrierIntrinsicID(std::optional<NVVM::BarrierReduction> reduction) {
-  if (reduction) {
-    switch (*reduction) {
-    case NVVM::BarrierReduction::AND:
-      return llvm::Intrinsic::nvvm_barrier0_and;
-    case NVVM::BarrierReduction::OR:
-      return llvm::Intrinsic::nvvm_barrier0_or;
-    case NVVM::BarrierReduction::POPC:
-      return llvm::Intrinsic::nvvm_barrier0_popc;
-    }
-  }
-
-  llvm_unreachable("Unknown reduction operation for barrier");
-}
-
 static unsigned getMembarIntrinsicID(NVVM::MemScopeKind scope) {
   switch (scope) {
   case NVVM::MemScopeKind::CTA: