Merge with nvvm.barrier

Author: clementval

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

@@ -977,54 +977,23 @@ def NVVM_Barrier0Op : NVVM_Op<"barrier0"> {
   }];
 }
 
-// Attrs describing the predicate of barrier0 operation.
-def Barrier0PredPopc : I32EnumAttrCase<"POPC", 0, "popc">;
-def Barrier0PredAnd : I32EnumAttrCase<"AND", 1, "and">;
-def Barrier0PredOr : I32EnumAttrCase<"OR", 2, "or">;
-
-def Barrier0Pred
-    : I32EnumAttr<"Barrier0Pred", "NVVM barrier0 predicate",
-                  [Barrier0PredPopc, Barrier0PredAnd, Barrier0PredOr]> {
+// Attrs describing the reduction operations for the barrier operation.
+def BarrierReductionPopc : I32EnumAttrCase<"POPC", 0, "popc">;
+def BarrierReductionAnd : I32EnumAttrCase<"AND", 1, "and">;
+def BarrierReductionOr : I32EnumAttrCase<"OR", 2, "or">;
+
+def BarrierReduction
+    : I32EnumAttr<"BarrierReduction", "NVVM barrier reduction operation",
+                  [BarrierReductionPopc, BarrierReductionAnd,
+                   BarrierReductionOr]> {
   let genSpecializedAttr = 0;
   let cppNamespace = "::mlir::NVVM";
 }
-def Barrier0PredAttr : EnumAttr<NVVM_Dialect, Barrier0Pred, "barrier0_pred"> {
+def BarrierReductionAttr
+    : EnumAttr<NVVM_Dialect, BarrierReduction, "reduction"> {
   let assemblyFormat = "`<` $value `>`";
 }
 
-def NVVM_Barrier0PredOp : NVVM_Op<"barrier0.pred">,
-                          Arguments<(ins Barrier0PredAttr:$pred, I32:$value)>,
-                          Results<(outs I32:$res)> {
-  let summary = "CTA Barrier Synchronization with predicate (Barrier ID 0)";
-  let description = [{
-    The `nvvm.barrier0` operation is a convenience operation that performs
-    barrier synchronization and communication within a CTA
-    (Cooperative Thread Array) using barrier ID 0. It is functionally
-    equivalent to `nvvm.barrier` or `nvvm.barrier id=0`.
-
-    `popc` is identical to `nvvm.barrier0` with the additional feature that it
-    evaluates predicate for all threads of the block and returns the number of
-    threads for which predicate evaluates to non-zero.
-
-    `and` is identical to `nvvm.barrier0` with the additional feature that it
-    evaluates predicate for all threads of the block and returns non-zero if
-    and only if predicate evaluates to non-zero for all of them.
-
-    `or` is identical to `nvvm.barrier0` with the additional feature that it
-    evaluates predicate for all threads of the block and returns non-zero if and
-    only if predicate evaluates to non-zero for any of them.
-
-    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-bar)
-  }];
-
-  let assemblyFormat =
-      "$value `:` type($value) $pred attr-dict `->` type($res)";
-  string llvmBuilder = [{
-      createIntrinsicCall(
-          builder, getBarrier0IntrinsicID($pred), {$value});
-  }];
-}
-
 def NVVM_BarrierOp : NVVM_Op<"barrier", [AttrSizedOperandSegments]> {
   let summary = "CTA Barrier Synchronization Op";
   let description = [{
@@ -1039,6 +1008,7 @@ def NVVM_BarrierOp : NVVM_Op<"barrier", [AttrSizedOperandSegments]> {
     - `numberOfThreads`: Specifies the number of threads participating in the barrier. 
       When specified, the value must be a multiple of the warp size. If not specified, 
       all threads in the CTA participate in the barrier.
+    - `reductionOp`
 
     The barrier operation guarantees that when the barrier completes, prior memory 
     accesses requested by participating threads are performed relative to all threads 
@@ -1055,31 +1025,36 @@ 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 arguments = (ins     
-    Optional<I32>:$barrierId,
-    Optional<I32>:$numberOfThreads);
+  let arguments = (ins Optional<I32>:$barrierId, Optional<I32>:$numberOfThreads,
+      OptionalAttr<BarrierReductionAttr>:$reductionOp,
+      Optional<I32>:$reductionPredicate);
   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), {$reductionPredicate});
     else
       createIntrinsicCall(
           builder, llvm::Intrinsic::nvvm_barrier_cta_sync_aligned_all, {id});
   }];
+  let results = (outs Optional<I32>:$res);
+
   let hasVerifier = 1;
 
-  let assemblyFormat = "(`id` `=` $barrierId^)? (`number_of_threads` `=` $numberOfThreads^)? attr-dict";
+  let assemblyFormat =
+      "(`id` `=` $barrierId^)? (`number_of_threads` `=` $numberOfThreads^)? "
+      "($reductionOp^ $reductionPredicate)? (`->` type($res)^)? attr-dict";
 
-  let builders = [
-    OpBuilder<(ins), [{
-      return build($_builder, $_state, Value{}, Value{});
+  let builders = [OpBuilder<(ins), [{
+      return build($_builder, $_state, TypeRange{}, Value{}, Value{}, {}, Value{});
     }]>,
-    OpBuilder<(ins "Value":$barrierId), [{
-      return build($_builder, $_state, barrierId, Value{});
-    }]>
-  ];
+                  OpBuilder<(ins "Value":$barrierId), [{
+      return build($_builder, $_state, TypeRange{}, barrierId, Value{}, {}, Value{});
+    }]>];
 }
 
 def NVVM_BarrierArriveOp : NVVM_PTXBuilder_Op<"barrier.arrive"> 

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

@@ -1504,6 +1504,15 @@ LogicalResult NVVM::BarrierOp::verify() {
   if (getNumberOfThreads() && !getBarrierId())
     return emitOpError(
         "barrier id is missing, it should be set between 0 to 15");
+
+  if (getBarrierId() && (getReductionOp() || getReductionPredicate()))
+    return emitOpError("reduction are only available for barrier id 0");
+
+  if ((getReductionOp() && !getReductionPredicate()) ||
+      (!getReductionOp() && getReductionPredicate()))
+    return emitOpError("reduction predicate and reduction operation must be "
+                       "specified together");
+
   return success();
 }
 

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

@@ -291,16 +291,20 @@ static unsigned getUnidirectionalFenceProxyID(NVVM::ProxyKind fromProxy,
   llvm_unreachable("Unsupported proxy kinds");
 }
 
-static unsigned getBarrier0IntrinsicID(NVVM::Barrier0Pred pred) {
-  switch (pred) {
-  case NVVM::Barrier0Pred::AND:
-    return llvm::Intrinsic::nvvm_barrier0_and;
-  case NVVM::Barrier0Pred::OR:
-    return llvm::Intrinsic::nvvm_barrier0_or;
-  case NVVM::Barrier0Pred::POPC:
-    return llvm::Intrinsic::nvvm_barrier0_popc;
+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 predicate for barrier0");
+
+  llvm_unreachable("Unknown reduction operation for barrier");
 }
 
 static unsigned getMembarIntrinsicID(NVVM::MemScopeKind scope) {

mlir/test/Target/LLVMIR/nvvm/barrier.mlir

@@ -0,0 +1,20 @@
+// RUN: mlir-translate -mlir-to-llvmir %s  -split-input-file --verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL: @llvm_nvvm_barrier(
+// CHECK-SAME: i32 %[[barId:.*]], i32 %[[numThreads:.*]], i32 %[[predicate:.*]])
+llvm.func @llvm_nvvm_barrier(%barID : i32, %numberOfThreads : i32, %predicate : i32) {
+  // CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 0)
+  nvvm.barrier
+  // CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 %[[barId]])
+  nvvm.barrier id = %barID
+  // CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.count(i32 %[[barId]], i32 %[[numThreads]])
+  nvvm.barrier id = %barID number_of_threads = %numberOfThreads
+  // CHECK: %{{.*}} = call i32 @llvm.nvvm.barrier0.and(i32 %[[predicate]])
+  %0 = nvvm.barrier #nvvm.reduction<and> %predicate -> i32
+  // CHECK: %{{.*}} = call i32 @llvm.nvvm.barrier0.or(i32 %[[predicate]])
+  %1 = nvvm.barrier #nvvm.reduction<or> %predicate -> i32
+  // CHECK: %{{.*}} = call i32 @llvm.nvvm.barrier0.popc(i32 %[[predicate]])
+  %2 = nvvm.barrier #nvvm.reduction<popc> %predicate -> i32
+
+  llvm.return
+}

mlir/test/Target/LLVMIR/nvvm/barrier0.mlir

@@ -166,18 +166,6 @@ llvm.func @nvvm_rcp(%0: f32) -> f32 {
   llvm.return %1 : f32
 }
 
-// CHECK-LABEL: @llvm_nvvm_barrier(
-// CHECK-SAME: i32 %[[barId:.*]], i32 %[[numThreads:.*]])
-llvm.func @llvm_nvvm_barrier(%barID : i32, %numberOfThreads : i32) {
-  // CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 0)
-  nvvm.barrier
-  // CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.all(i32 %[[barId]])
-  nvvm.barrier id = %barID
-  // CHECK: call void @llvm.nvvm.barrier.cta.sync.aligned.count(i32 %[[barId]], i32 %[[numThreads]])
-  nvvm.barrier id = %barID number_of_threads = %numberOfThreads
-  llvm.return
-}
-
 // CHECK-LABEL: @llvm_nvvm_cluster_arrive
 llvm.func @llvm_nvvm_cluster_arrive() {
   // CHECK: call void @llvm.nvvm.barrier.cluster.arrive()