[Mlir-commits] [mlir] [mlir][acc] Add acc.on_device op (PR #208096)

[Ivan R. Ivanov]

https://github.com/ivanradanov updated https://github.com/llvm/llvm-project/pull/208096

>From fd41d92a901dabf805d1117b7263a6f74f3727dd Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <iivanov at nvidia.com>
Date: Tue, 7 Jul 2026 12:25:07 -0700
Subject: [PATCH 1/2] [mlir][acc] Add acc.on_device op

Add an operation to represent the runtime call to acc_on_device.

This runtime call is important to fold early in the compilation pipeline
and having an operation allows us to easily recognize it when emitted by
frontends.
---
 .../mlir/Dialect/OpenACC/OpenACCOps.td        | 27 +++++++++++++++++++
 mlir/test/Dialect/OpenACC/ops.mlir            |  8 ++++++
 2 files changed, 35 insertions(+)

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index 26a9344fa3a15..8af352c08c156 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
@@ -3713,6 +3713,33 @@ def OpenACC_WaitOp : OpenACC_Op<"wait", [AttrSizedOperandSegments]> {
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// acc_on_device runtime call
+//===----------------------------------------------------------------------===//
+
+def OpenACC_OnDeviceOp : OpenACC_Op<"on_device", [NoMemoryEffect]> {
+  let summary = "acc_on_device intrinsic operation";
+
+  let description = [{
+    Represents a call to the OpenACC `acc_on_device` runtime function.
+    Returns whether the current thread is executing on the given device type.
+
+    Example:
+
+    ```mlir
+    %host = arith.constant 1 : i32
+    %on_host = acc.on_device %host : i32 -> i1
+    ```
+  }];
+
+  let arguments = (ins I32:$deviceType);
+  let results = (outs I1:$result);
+
+  let assemblyFormat = [{
+    $deviceType attr-dict `:` type($deviceType) `->` type($result)
+  }];
+}
+
 include "mlir/Dialect/OpenACC/OpenACCCGAttributes.td"
 include "mlir/Dialect/OpenACC/OpenACCCGOps.td"
 
diff --git a/mlir/test/Dialect/OpenACC/ops.mlir b/mlir/test/Dialect/OpenACC/ops.mlir
index 2fb73e400001f..022b71071950e 100644
--- a/mlir/test/Dialect/OpenACC/ops.mlir
+++ b/mlir/test/Dialect/OpenACC/ops.mlir
@@ -1187,6 +1187,14 @@ acc.shutdown device_num(%idxValue : index) if(%ifCond)
 
 // -----
 
+%hostDevType = arith.constant 2 : i32
+%onHost = acc.on_device %hostDevType : i32 -> i1
+
+// CHECK: [[HOSTDEVTYPE:%.*]] = arith.constant 2 : i32
+// CHECK: %[[ONHOST:.*]] = acc.on_device [[HOSTDEVTYPE]] : i32 -> i1
+
+// -----
+
 func.func @testexitdataop(%a: !llvm.ptr) -> () {
   %ifCond = arith.constant true
   %i64Value = arith.constant 1 : i64

>From 7927ec2fc15f204cc4f14cd3cf1f825f9cd39fa3 Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <iivanov at nvidia.com>
Date: Wed, 8 Jul 2026 07:54:40 -0700
Subject: [PATCH 2/2] op fix

---
 mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td |  4 ++--
 mlir/test/Dialect/OpenACC/ops.mlir              | 14 +++++++++-----
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index 8af352c08c156..084536dd79960 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
@@ -3717,7 +3717,7 @@ def OpenACC_WaitOp : OpenACC_Op<"wait", [AttrSizedOperandSegments]> {
 // acc_on_device runtime call
 //===----------------------------------------------------------------------===//
 
-def OpenACC_OnDeviceOp : OpenACC_Op<"on_device", [NoMemoryEffect]> {
+def OpenACC_OnDeviceOp : OpenACC_Op<"on_device"> {
   let summary = "acc_on_device intrinsic operation";
 
   let description = [{
@@ -3732,7 +3732,7 @@ def OpenACC_OnDeviceOp : OpenACC_Op<"on_device", [NoMemoryEffect]> {
     ```
   }];
 
-  let arguments = (ins I32:$deviceType);
+  let arguments = (ins IntOrIndex:$deviceType);
   let results = (outs I1:$result);
 
   let assemblyFormat = [{
diff --git a/mlir/test/Dialect/OpenACC/ops.mlir b/mlir/test/Dialect/OpenACC/ops.mlir
index 022b71071950e..c696cefb94600 100644
--- a/mlir/test/Dialect/OpenACC/ops.mlir
+++ b/mlir/test/Dialect/OpenACC/ops.mlir
@@ -1187,11 +1187,15 @@ acc.shutdown device_num(%idxValue : index) if(%ifCond)
 
 // -----
 
-%hostDevType = arith.constant 2 : i32
-%onHost = acc.on_device %hostDevType : i32 -> i1
-
-// CHECK: [[HOSTDEVTYPE:%.*]] = arith.constant 2 : i32
-// CHECK: %[[ONHOST:.*]] = acc.on_device [[HOSTDEVTYPE]] : i32 -> i1
+%devTypeInt = arith.constant 2 : i32
+%onHostInt = acc.on_device %devType : i32 -> i1
+%devTypeIndex = arith.constant 2 : index
+%onHostIndex = acc.on_device %devType : index -> i1
+
+// CHECK: [[DEVTYPEINT:%.*]] = arith.constant 2 : i32
+// CHECK: %[[ONHOSTINT:.*]] = acc.on_device [[DEVTYPEINT]] : i32 -> i1
+// CHECK: [[DEVTYPEINDEX:%.*]] = arith.constant 2 : index
+// CHECK: %[[ONHOSTINDEX:.*]] = acc.on_device [[DEVTYPEINDEX]] : index -> i1
 
 // -----