Reapply "[MLIR][Python] add ctype python binding support for bf16" #101271
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…lvm#93771) This reverts commit e6821dd.
xurui1995
requested review from
ftynse,
makslevental and
stellaraccident
as code owners
|
@llvm/pr-subscribers-mlir Author: Bimo (xurui1995) ChangesReapply the PR which was reverted due to built-bots, and now the bots get updated. Full diff: https://github.com/llvm/llvm-project/pull/101271.diff 3 Files Affected:
diff --git a/mlir/python/mlir/runtime/np_to_memref.py b/mlir/python/mlir/runtime/np_to_memref.py
index f6b706f9bc8ae..882b2751921bf 100644
--- a/mlir/python/mlir/runtime/np_to_memref.py
+++ b/mlir/python/mlir/runtime/np_to_memref.py
@@ -7,6 +7,12 @@
import numpy as np
import ctypes
+try:
+ import ml_dtypes
+except ModuleNotFoundError:
+ # The third-party ml_dtypes provides some optional low precision data-types for NumPy.
+ ml_dtypes = None
+
class C128(ctypes.Structure):
"""A ctype representation for MLIR's Double Complex."""
@@ -26,6 +32,12 @@ class F16(ctypes.Structure):
_fields_ = [("f16", ctypes.c_int16)]
+class BF16(ctypes.Structure):
+ """A ctype representation for MLIR's BFloat16."""
+
+ _fields_ = [("bf16", ctypes.c_int16)]
+
+
# https://stackoverflow.com/questions/26921836/correct-way-to-test-for-numpy-dtype
def as_ctype(dtp):
"""Converts dtype to ctype."""
@@ -35,6 +47,8 @@ def as_ctype(dtp):
return C64
if dtp == np.dtype(np.float16):
return F16
+ if ml_dtypes is not None and dtp == ml_dtypes.bfloat16:
+ return BF16
return np.ctypeslib.as_ctypes_type(dtp)
@@ -46,6 +60,11 @@ def to_numpy(array):
return array.view("complex64")
if array.dtype == F16:
return array.view("float16")
+ assert not (
+ array.dtype == BF16 and ml_dtypes is None
+ ), f"bfloat16 requires the ml_dtypes package, please run:\n\npip install ml_dtypes\n"
+ if array.dtype == BF16:
+ return array.view("bfloat16")
return array
diff --git a/mlir/python/requirements.txt b/mlir/python/requirements.txt
index acd6dbb25edaf..6ec63e43adf89 100644
--- a/mlir/python/requirements.txt
+++ b/mlir/python/requirements.txt
@@ -1,3 +1,4 @@
numpy>=1.19.5, <=1.26
pybind11>=2.9.0, <=2.10.3
-PyYAML>=5.3.1, <=6.0.1
\ No newline at end of file
+PyYAML>=5.3.1, <=6.0.1
+ml_dtypes # provides several NumPy dtype extensions, including the bf16
\ No newline at end of file
diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py
index e8b47007a8907..8125bf3fb8fc9 100644
--- a/mlir/test/python/execution_engine.py
+++ b/mlir/test/python/execution_engine.py
@@ -5,6 +5,7 @@
from mlir.passmanager import *
from mlir.execution_engine import *
from mlir.runtime import *
+from ml_dtypes import bfloat16
# Log everything to stderr and flush so that we have a unified stream to match
@@ -521,6 +522,45 @@ def testComplexUnrankedMemrefAdd():
run(testComplexUnrankedMemrefAdd)
+# Test bf16 memrefs
+# CHECK-LABEL: TEST: testBF16Memref
+def testBF16Memref():
+ with Context():
+ module = Module.parse(
+ """
+ module {
+ func.func @main(%arg0: memref<1xbf16>,
+ %arg1: memref<1xbf16>) attributes { llvm.emit_c_interface } {
+ %0 = arith.constant 0 : index
+ %1 = memref.load %arg0[%0] : memref<1xbf16>
+ memref.store %1, %arg1[%0] : memref<1xbf16>
+ return
+ }
+ } """
+ )
+
+ arg1 = np.array([0.5]).astype(bfloat16)
+ arg2 = np.array([0.0]).astype(bfloat16)
+
+ arg1_memref_ptr = ctypes.pointer(
+ ctypes.pointer(get_ranked_memref_descriptor(arg1))
+ )
+ arg2_memref_ptr = ctypes.pointer(
+ ctypes.pointer(get_ranked_memref_descriptor(arg2))
+ )
+
+ execution_engine = ExecutionEngine(lowerToLLVM(module))
+ execution_engine.invoke("main", arg1_memref_ptr, arg2_memref_ptr)
+
+ # test to-numpy utility
+ # CHECK: [0.5]
+ npout = ranked_memref_to_numpy(arg2_memref_ptr[0])
+ log(npout)
+
+
+run(testBF16Memref)
+
+
# Test addition of two 2d_memref
# CHECK-LABEL: TEST: testDynamicMemrefAdd2D
def testDynamicMemrefAdd2D():
|
|
@joker-eph, hi, could you help approve this? |
joker-eph
approved these changes
Jul 31, 2024
|
I'll merge it now and keep an eye on the bot in case. |
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/177/builds/2495 Here is the relevant piece of the build log for the reference: |
same problem in running mlir-rocm-mi200 |
Previous
Now
|
5 tasks
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Reapply the PR which was reverted due to built-bots, and now the bots get updated.
https://discourse.llvm.org/t/need-a-help-with-the-built-bots/79437
original PR: #92489, reverted in #93771