[Mlir-commits] [mlir] 8481fb1 - [MLIR][ROCDL] Fix BallotOp LLVM translation and add doc (#85116)

[llvmlistbot at llvm.org]

Author: Zahi Moudallal
Date: 2024-03-14T08:43:48-07:00
New Revision: 8481fb1698a4a54b3965a1654046df736a57e144

URL: https://github.com/llvm/llvm-project/commit/8481fb1698a4a54b3965a1654046df736a57e144
DIFF: https://github.com/llvm/llvm-project/commit/8481fb1698a4a54b3965a1654046df736a57e144.diff

LOG: [MLIR][ROCDL] Fix BallotOp LLVM translation and add doc (#85116)

This modifies the return type of the intrinsic call to handle 32 and 64
bits
properly and document the MLIR operation.

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
    mlir/test/Target/LLVMIR/rocdl.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index abb38a3df8068c..1dabf5d7979b70 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -162,10 +162,18 @@ def ROCDL_BallotOp :
   ROCDL_Op<"ballot">,
   Results<(outs LLVM_Type:$res)>,
   Arguments<(ins I1:$pred)> {
+  let summary = "Vote across thread group";
+
+  let description = [{
+      Ballot provides a bit mask containing the 1-bit predicate value from each lane. 
+      The nth bit of the result contains the 1 bit contributed by the nth warp lane.
+  }];
+
   string llvmBuilder = [{
       $res = createIntrinsicCall(builder,
-            llvm::Intrinsic::amdgcn_ballot, {$pred}, {llvm::Type::getInt32Ty(moduleTranslation.getLLVMContext())});
+            llvm::Intrinsic::amdgcn_ballot, {$pred}, {$_resultType});
   }];
+
   let assemblyFormat = "$pred attr-dict `:` type($res)";
 }
 

diff  --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index 93550f5c7bd5a4..ce6b56d48437a0 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -88,13 +88,20 @@ llvm.func @rocdl.bpermute(%src : i32) -> i32 {
   llvm.return %0 : i32
 }
 
-llvm.func @rocdl.ballot(%pred : i1) -> i32 {
-  // CHECK-LABEL: rocdl.ballot
+llvm.func @rocdl.ballot32(%pred : i1) -> i32 {
+  // CHECK-LABEL: rocdl.ballot32
   // CHECK: call i32 @llvm.amdgcn.ballot
   %0 = rocdl.ballot %pred : i32
   llvm.return %0 : i32
 }
 
+llvm.func @rocdl.ballot64(%pred : i1) -> i64 {
+  // CHECK-LABEL: rocdl.ballot64
+  // CHECK: call i64 @llvm.amdgcn.ballot
+  %0 = rocdl.ballot %pred : i64
+  llvm.return %0 : i64
+}
+
 llvm.func @rocdl.waitcnt() {
   // CHECK-LABEL: rocdl.waitcnt
   // CHECK-NEXT: call void @llvm.amdgcn.s.waitcnt(i32 0)