[Mlir-commits] [mlir] [MLIR][ROCDL] Added summary and description to BallotOp (PR #85116)

[Zahi Moudallal]

https://github.com/zahimoud updated https://github.com/llvm/llvm-project/pull/85116

>From 34130f377ebf243f9976517fc8c40d47059f6d31 Mon Sep 17 00:00:00 2001
From: Zahi Moudallal <zahi at openai.com>
Date: Wed, 13 Mar 2024 11:37:02 -0700
Subject: [PATCH 1/3] [MLIR][ROCDL] Added summary and description to BallotOp

---
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index abb38a3df8068c..8c641b931de18e 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -162,10 +162,20 @@ def ROCDL_BallotOp :
   ROCDL_Op<"ballot">,
   Results<(outs LLVM_Type:$res)>,
   Arguments<(ins I1:$pred)> {
+  let summary = "Vote across thread group";
+
+  let description = [{
+      This operation applies a reduction to the source predicate across all the active threads within a warp, 
+      resulting in a destination predicate value that is uniform for every thread in the warp.
+
+      i1 -> i32
+  }];
+
   string llvmBuilder = [{
       $res = createIntrinsicCall(builder,
             llvm::Intrinsic::amdgcn_ballot, {$pred}, {llvm::Type::getInt32Ty(moduleTranslation.getLLVMContext())});
   }];
+
   let assemblyFormat = "$pred attr-dict `:` type($res)";
 }
 

>From 31632cc460715377642f519512ddb05298d4bf8f Mon Sep 17 00:00:00 2001
From: Zahi Moudallal <zahi at openai.com>
Date: Wed, 13 Mar 2024 11:49:53 -0700
Subject: [PATCH 2/3] [MLIR][ROCDL] Fixed documentation and return type.

---
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 8c641b931de18e..cfa366c3146518 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -165,15 +165,15 @@ def ROCDL_BallotOp :
   let summary = "Vote across thread group";
 
   let description = [{
-      This operation applies a reduction to the source predicate across all the active threads within a warp, 
-      resulting in a destination predicate value that is uniform for every thread in the warp.
+      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
 
-      i1 -> i32
+      i1 -> anyint
   }];
 
   string llvmBuilder = [{
       $res = createIntrinsicCall(builder,
-            llvm::Intrinsic::amdgcn_ballot, {$pred}, {llvm::Type::getInt32Ty(moduleTranslation.getLLVMContext())});
+            llvm::Intrinsic::amdgcn_ballot, {$pred}, {$res.getType()});
   }];
 
   let assemblyFormat = "$pred attr-dict `:` type($res)";

>From a27198af247d9d74187e8c1a17afe3c94ecbb21e Mon Sep 17 00:00:00 2001
From: Zahi Moudallal <zahi at openai.com>
Date: Wed, 13 Mar 2024 11:58:27 -0700
Subject: [PATCH 3/3] [MLIR][ROCDL] Fixed documentation and added lit test

---
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 4 +---
 mlir/test/Target/LLVMIR/rocdl.mlir           | 7 +++++++
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index cfa366c3146518..cd92d64e5f66b5 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -166,9 +166,7 @@ def ROCDL_BallotOp :
 
   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
-
-      i1 -> anyint
+      The nth bit of the result contains the 1 bit contributed by the nth warp lane.
   }];
 
   string llvmBuilder = [{
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index 93550f5c7bd5a4..7b48ac4239b5a2 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -95,6 +95,13 @@ llvm.func @rocdl.ballot(%pred : i1) -> i32 {
   llvm.return %0 : i32
 }
 
+llvm.func @rocdl.ballot(%pred : i1) -> i64 {
+  // CHECK-LABEL: rocdl.ballot
+  // 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)