[Mlir-commits] [mlir] a7d80c5 - [MLIR][python bindings] add vendor gpu dialects

llvmlistbot at llvm.org llvmlistbot at llvm.org
Sun Aug 13 15:05:48 PDT 2023 

Author: max
Date: 2023-08-13T16:45:20-05:00
New Revision: a7d80c50aa9121bff08d59ab4c4d7210dd2400f9

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

LOG: [MLIR][python bindings] add vendor gpu dialects

Differential Revision: https://reviews.llvm.org/D157820

Added: 
    mlir/include/mlir-c/Dialect/AMDGPU.h
    mlir/include/mlir-c/Dialect/NVGPU.h
    mlir/include/mlir-c/Dialect/NVVM.h
    mlir/include/mlir-c/Dialect/ROCDL.h
    mlir/lib/CAPI/Dialect/AMDGPU.cpp
    mlir/lib/CAPI/Dialect/NVGPU.cpp
    mlir/lib/CAPI/Dialect/NVVM.cpp
    mlir/lib/CAPI/Dialect/ROCDL.cpp
    mlir/python/mlir/dialects/AMDGPUOps.td
    mlir/python/mlir/dialects/NVGPUOps.td
    mlir/python/mlir/dialects/NVVMOps.td
    mlir/python/mlir/dialects/ROCDLOps.td
    mlir/python/mlir/dialects/amdgpu.py
    mlir/python/mlir/dialects/nvgpu.py
    mlir/python/mlir/dialects/nvvm.py
    mlir/python/mlir/dialects/rocdl.py
    mlir/test/python/dialects/amdgpu.py
    mlir/test/python/dialects/nvgpu.py
    mlir/test/python/dialects/nvvm.py
    mlir/test/python/dialects/rocdl.py

Modified: 
    mlir/lib/CAPI/Dialect/CMakeLists.txt
    mlir/python/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir-c/Dialect/AMDGPU.h b/mlir/include/mlir-c/Dialect/AMDGPU.h
new file mode 100644
index 00000000000000..142044f7f3afe2
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/AMDGPU.h
@@ -0,0 +1,25 @@
+//===-- mlir-c/Dialect/AMDGPU.h - C API for AMDGPU dialect --*- C -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_AMDGPU_H
+#define MLIR_C_DIALECT_AMDGPU_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(AMDGPU, amdgpu);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_AMDGPU_H

diff  --git a/mlir/include/mlir-c/Dialect/NVGPU.h b/mlir/include/mlir-c/Dialect/NVGPU.h
new file mode 100644
index 00000000000000..580d566794c09f
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/NVGPU.h
@@ -0,0 +1,25 @@
+//===-- mlir-c/Dialect/NVGPU.h - C API for NVGPU dialect --*- C -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_NVGPU_H
+#define MLIR_C_DIALECT_NVGPU_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(NVGPU, nvgpu);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_NVGPU_H

diff  --git a/mlir/include/mlir-c/Dialect/NVVM.h b/mlir/include/mlir-c/Dialect/NVVM.h
new file mode 100644
index 00000000000000..cf5d9301d4b284
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/NVVM.h
@@ -0,0 +1,25 @@
+//===-- mlir-c/Dialect/NVVM.h - C API for NVVM dialect --*- C -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_NVVM_H
+#define MLIR_C_DIALECT_NVVM_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(NVVM, nvvm);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_NVVM_H

diff  --git a/mlir/include/mlir-c/Dialect/ROCDL.h b/mlir/include/mlir-c/Dialect/ROCDL.h
new file mode 100644
index 00000000000000..e5dbb55b581878
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/ROCDL.h
@@ -0,0 +1,25 @@
+//===-- mlir-c/Dialect/ROCDL.h - C API for ROCDL dialect --*- C -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_ROCDL_H
+#define MLIR_C_DIALECT_ROCDL_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(ROCDL, rocdl);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_ROCDL_H

diff  --git a/mlir/lib/CAPI/Dialect/AMDGPU.cpp b/mlir/lib/CAPI/Dialect/AMDGPU.cpp
new file mode 100644
index 00000000000000..28efe602511ea6
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/AMDGPU.cpp
@@ -0,0 +1,14 @@
+//===- AMDGPU.cpp - C Interface for AMDGPU dialect ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/AMDGPU.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(AMDGPU, ml_program,
+                                      mlir::amdgpu::AMDGPUDialect)

diff  --git a/mlir/lib/CAPI/Dialect/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt
index 4b4ab74e6fc341..13e57a29d65c92 100644
--- a/mlir/lib/CAPI/Dialect/CMakeLists.txt
+++ b/mlir/lib/CAPI/Dialect/CMakeLists.txt
@@ -1,3 +1,12 @@
+add_mlir_upstream_c_api_library(MLIRCAPIAMDGPU
+  AMDGPU.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRAMDGPUDialect
+)
+
 add_mlir_upstream_c_api_library(MLIRCAPIArith
   Arith.cpp
 
@@ -96,6 +105,34 @@ add_mlir_upstream_c_api_library(MLIRCAPIMLProgram
   MLIRMLProgramDialect
 )
 
+add_mlir_upstream_c_api_library(MLIRCAPINVGPU
+  NVGPU.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRNVGPUDialect
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPINVVM
+  NVVM.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRNVVMDialect
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIROCDL
+  ROCDL.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRROCDLDialect
+)
+
+
 add_mlir_upstream_c_api_library(MLIRCAPISCF
   SCF.cpp
 

diff  --git a/mlir/lib/CAPI/Dialect/NVGPU.cpp b/mlir/lib/CAPI/Dialect/NVGPU.cpp
new file mode 100644
index 00000000000000..02d10954a03776
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/NVGPU.cpp
@@ -0,0 +1,13 @@
+//===- NVGPU.cpp - C Interface for NVGPU dialect ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/NVGPU.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(NVGPU, nvgpu, mlir::nvgpu::NVGPUDialect)

diff  --git a/mlir/lib/CAPI/Dialect/NVVM.cpp b/mlir/lib/CAPI/Dialect/NVVM.cpp
new file mode 100644
index 00000000000000..a87581664bbaba
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/NVVM.cpp
@@ -0,0 +1,13 @@
+//===- NVVM.cpp - C Interface for NVVM dialect ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/NVVM.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(NVVM, nvvm, mlir::NVVM::NVVMDialect)

diff  --git a/mlir/lib/CAPI/Dialect/ROCDL.cpp b/mlir/lib/CAPI/Dialect/ROCDL.cpp
new file mode 100644
index 00000000000000..63e2fa881670d8
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/ROCDL.cpp
@@ -0,0 +1,13 @@
+//===- ROCDL.cpp - C Interface for ROCDL dialect ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/ROCDL.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(ROCDL, rocdl, mlir::ROCDL::ROCDLDialect)

diff  --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 0fae5dbb86347a..05d09eaf7b8b25 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -46,6 +46,14 @@ declare_mlir_python_sources(MLIRPythonCAPI.HeaderSources
 # Dialect bindings
 ################################################################################
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/AMDGPUOps.td
+  SOURCES
+    dialects/amdgpu.py
+  DIALECT_NAME amdgpu)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
@@ -264,6 +272,30 @@ declare_mlir_dialect_python_bindings(
     dialects/_ml_program_ops_ext.py
   DIALECT_NAME ml_program)
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/NVGPUOps.td
+  SOURCES
+    dialects/nvgpu.py
+  DIALECT_NAME nvgpu)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/NVVMOps.td
+  SOURCES
+    dialects/nvvm.py
+  DIALECT_NAME nvvm)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/ROCDLOps.td
+  SOURCES
+    dialects/rocdl.py
+  DIALECT_NAME rocdl)
+
 declare_mlir_python_sources(
   MLIRPythonSources.Dialects.quant
   ADD_TO_PARENT MLIRPythonSources.Dialects

diff  --git a/mlir/python/mlir/dialects/AMDGPUOps.td b/mlir/python/mlir/dialects/AMDGPUOps.td
new file mode 100644
index 00000000000000..fe9371971e105b
--- /dev/null
+++ b/mlir/python/mlir/dialects/AMDGPUOps.td
@@ -0,0 +1,14 @@
+//===-- AMDGPUOps.td - Entry point for AMDGPUOps -----*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_AMDGPU_OPS
+#define PYTHON_BINDINGS_AMDGPU_OPS
+
+include "mlir/Dialect/AMDGPU/IR/AMDGPU.td"
+
+#endif

diff  --git a/mlir/python/mlir/dialects/NVGPUOps.td b/mlir/python/mlir/dialects/NVGPUOps.td
new file mode 100644
index 00000000000000..ae54822cd90708
--- /dev/null
+++ b/mlir/python/mlir/dialects/NVGPUOps.td
@@ -0,0 +1,14 @@
+//===-- NVGPUOps.td - Entry point for NVGPUOps -----*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_NVGPU_OPS
+#define PYTHON_BINDINGS_NVGPU_OPS
+
+include "mlir/Dialect/NVGPU/IR/NVGPU.td"
+
+#endif

diff  --git a/mlir/python/mlir/dialects/NVVMOps.td b/mlir/python/mlir/dialects/NVVMOps.td
new file mode 100644
index 00000000000000..f57d204a898e88
--- /dev/null
+++ b/mlir/python/mlir/dialects/NVVMOps.td
@@ -0,0 +1,14 @@
+//===-- NVVMOps.td - Entry point for NVVMOps -----*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_NVVM_OPS
+#define PYTHON_BINDINGS_NVVM_OPS
+
+include "mlir/Dialect/LLVMIR/NVVMOps.td"
+
+#endif

diff  --git a/mlir/python/mlir/dialects/ROCDLOps.td b/mlir/python/mlir/dialects/ROCDLOps.td
new file mode 100644
index 00000000000000..fa5c9ebc326567
--- /dev/null
+++ b/mlir/python/mlir/dialects/ROCDLOps.td
@@ -0,0 +1,14 @@
+//===-- ROCDLOps.td - Entry point for ROCDLOps -----*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_ROCDL_OPS
+#define PYTHON_BINDINGS_ROCDL_OPS
+
+include "mlir/Dialect/LLVMIR/ROCDLOps.td"
+
+#endif

diff  --git a/mlir/python/mlir/dialects/amdgpu.py b/mlir/python/mlir/dialects/amdgpu.py
new file mode 100644
index 00000000000000..35283278e8fb03
--- /dev/null
+++ b/mlir/python/mlir/dialects/amdgpu.py
@@ -0,0 +1,5 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._amdgpu_ops_gen import *

diff  --git a/mlir/python/mlir/dialects/nvgpu.py b/mlir/python/mlir/dialects/nvgpu.py
new file mode 100644
index 00000000000000..afd570cae5300f
--- /dev/null
+++ b/mlir/python/mlir/dialects/nvgpu.py
@@ -0,0 +1,5 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._nvgpu_ops_gen import *

diff  --git a/mlir/python/mlir/dialects/nvvm.py b/mlir/python/mlir/dialects/nvvm.py
new file mode 100644
index 00000000000000..87b2a4fd6bf853
--- /dev/null
+++ b/mlir/python/mlir/dialects/nvvm.py
@@ -0,0 +1,5 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._nvvm_ops_gen import *

diff  --git a/mlir/python/mlir/dialects/rocdl.py b/mlir/python/mlir/dialects/rocdl.py
new file mode 100644
index 00000000000000..aa47cb4b55792b
--- /dev/null
+++ b/mlir/python/mlir/dialects/rocdl.py
@@ -0,0 +1,5 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._rocdl_ops_gen import *

diff  --git a/mlir/test/python/dialects/amdgpu.py b/mlir/test/python/dialects/amdgpu.py
new file mode 100644
index 00000000000000..c8039d494cf814
--- /dev/null
+++ b/mlir/test/python/dialects/amdgpu.py
@@ -0,0 +1,22 @@
+# RUN: %PYTHON %s | FileCheck %s
+# This is just a smoke test that the dialect is functional.
+
+from mlir.ir import *
+from mlir.dialects import amdgpu, arith, memref
+
+
+def constructAndPrintInModule(f):
+    print("\nTEST:", f.__name__)
+    with Context(), Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            f()
+        print(module)
+    return f
+
+
+# CHECK-LABEL: testSmoke
+ at constructAndPrintInModule
+def testSmoke():
+    # CHECK: amdgpu.lds_barrier
+    amdgpu.LDSBarrierOp()

diff  --git a/mlir/test/python/dialects/nvgpu.py b/mlir/test/python/dialects/nvgpu.py
new file mode 100644
index 00000000000000..3158388f0e6869
--- /dev/null
+++ b/mlir/test/python/dialects/nvgpu.py
@@ -0,0 +1,26 @@
+# RUN: %PYTHON %s | FileCheck %s
+# This is just a smoke test that the dialect is functional.
+
+from mlir.ir import *
+from mlir.dialects import nvgpu, arith, memref
+
+
+def constructAndPrintInModule(f):
+    print("\nTEST:", f.__name__)
+    with Context(), Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            f()
+        print(module)
+    return f
+
+
+# CHECK-LABEL: testSmoke
+ at constructAndPrintInModule
+def testSmoke():
+    cst = arith.ConstantOp(value=42, result=IndexType.get())
+    mem_t = MemRefType.get((10, 10), F32Type.get(), memory_space=Attribute.parse("3"))
+    vec_t = VectorType.get((4, 1), F32Type.get())
+    mem = memref.AllocOp(mem_t, [], [])
+    # CHECK: %0 = nvgpu.ldmatrix %alloc[%c42, %c42] {numTiles = 4 : i32, transpose = false} : memref<10x10xf32, 3> -> vector<4x1xf32>
+    nvgpu.LdMatrixOp(vec_t, mem, [cst, cst], False, 4)

diff  --git a/mlir/test/python/dialects/nvvm.py b/mlir/test/python/dialects/nvvm.py
new file mode 100644
index 00000000000000..7d68a151345ca2
--- /dev/null
+++ b/mlir/test/python/dialects/nvvm.py
@@ -0,0 +1,22 @@
+# RUN: %PYTHON %s | FileCheck %s
+# This is just a smoke test that the dialect is functional.
+
+from mlir.ir import *
+from mlir.dialects import nvvm
+
+
+def constructAndPrintInModule(f):
+    print("\nTEST:", f.__name__)
+    with Context(), Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            f()
+        print(module)
+    return f
+
+
+# CHECK-LABEL: testSmoke
+ at constructAndPrintInModule
+def testSmoke():
+    # CHECK: nvvm.cp.async.wait.group 5
+    nvvm.CpAsyncWaitGroupOp(5)

diff  --git a/mlir/test/python/dialects/rocdl.py b/mlir/test/python/dialects/rocdl.py
new file mode 100644
index 00000000000000..a4eca2766899ba
--- /dev/null
+++ b/mlir/test/python/dialects/rocdl.py
@@ -0,0 +1,22 @@
+# RUN: %PYTHON %s | FileCheck %s
+# This is just a smoke test that the dialect is functional.
+
+from mlir.ir import *
+from mlir.dialects import rocdl
+
+
+def constructAndPrintInModule(f):
+    print("\nTEST:", f.__name__)
+    with Context(), Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            f()
+        print(module)
+    return f
+
+
+# CHECK-LABEL: testSmoke
+ at constructAndPrintInModule
+def testSmoke():
+    # CHECK: rocdl.barrier
+    rocdl.BarrierOp()

More information about the Mlir-commits mailing list