[Mlir-commits] [mlir] [MLIR][ROCDL] Added BallotOp and lit test (PR #84856)

[Zahi Moudallal]

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

>From 6e8a877168363cb36196b9a389bbd4336da14066 Mon Sep 17 00:00:00 2001
From: Zahi Moudallal <zahi at openai.com>
Date: Mon, 11 Mar 2024 17:06:23 -0700
Subject: [PATCH 1/2] [ROCDL] Added BallotOp and lit test

---
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 10 ++++++++++
 mlir/test/Target/LLVMIR/rocdl.mlir           |  7 +++++++
 2 files changed, 17 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 32b5a1c016b6f8..600d4c943adf7e 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -158,6 +158,16 @@ Arguments<(ins I32:$index,
    }];
 }
 
+def ROCDL_BallotOp :
+  ROCDL_Op<"ballot">,
+  Results<(outs LLVM_Type:$res)>,
+  Arguments<(ins I1:$pred)> {
+  string llvmBuilder = [{
+      $res = createIntrinsicCall(builder,
+            llvm::Intrinsic::amdgcn_ballot, {$pred});
+  }];
+  let assemblyFormat = "$pred attr-dict `:` type($res)";
+}
 
 //===----------------------------------------------------------------------===//
 // Thread index and Block index
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index d35acb0475e6f8..93550f5c7bd5a4 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -88,6 +88,13 @@ llvm.func @rocdl.bpermute(%src : i32) -> i32 {
   llvm.return %0 : i32
 }
 
+llvm.func @rocdl.ballot(%pred : i1) -> i32 {
+  // CHECK-LABEL: rocdl.ballot
+  // CHECK: call i32 @llvm.amdgcn.ballot
+  %0 = rocdl.ballot %pred : i32
+  llvm.return %0 : i32
+}
+
 llvm.func @rocdl.waitcnt() {
   // CHECK-LABEL: rocdl.waitcnt
   // CHECK-NEXT: call void @llvm.amdgcn.s.waitcnt(i32 0)

>From 1ea220b007a45e70597a7c79f0ce734f259d3658 Mon Sep 17 00:00:00 2001
From: Zahi Moudallal <zahi at openai.com>
Date: Tue, 12 Mar 2024 11:18:28 -0700
Subject: [PATCH 2/2] Attempt to add result type to intrinsic call

---
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 600d4c943adf7e..abb38a3df8068c 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -164,7 +164,7 @@ def ROCDL_BallotOp :
   Arguments<(ins I1:$pred)> {
   string llvmBuilder = [{
       $res = createIntrinsicCall(builder,
-            llvm::Intrinsic::amdgcn_ballot, {$pred});
+            llvm::Intrinsic::amdgcn_ballot, {$pred}, {llvm::Type::getInt32Ty(moduleTranslation.getLLVMContext())});
   }];
   let assemblyFormat = "$pred attr-dict `:` type($res)";
 }