Skip to content

Commit 0703e70

Browse files
pranavm-nvidiapranavm-nvidia
committed
[mlir][python] Add bindings for mlirDenseElementsAttrGet
This change adds bindings for `mlirDenseElementsAttrGet` which accepts a list of MLIR attributes and constructs a DenseElementsAttr. This allows for creating `DenseElementsAttr`s of types not natively supported by Python (e.g. BF16) without requiring other dependencies (e.g. `numpy` + `ml-dtypes`).
1 parent 9f23138 commit 0703e70

File tree

2 files changed

+161
-5
lines changed

2 files changed

+161
-5
lines changed

‎mlir/lib/Bindings/Python/IRAttributes.cpp

Lines changed: 79 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,26 @@ or 255), then a splat will be created.
7272
type or if the buffer does not meet expectations.
7373
)";
7474

75+
static const char kDenseElementsAttrGetFromListDocstring[] =
76+
R"(Gets a DenseElementsAttr from a Python list of attributes.
77+
78+
Note that it can be expensive to construct attributes individually.
79+
For a large number of elements, consider using a Python buffer or array instead.
80+
81+
Args:
82+
attrs: A list of attributes.
83+
type: The desired shape and type of the resulting DenseElementsAttr.
84+
If not provided, the element type is determined based on the type
85+
of the 0th attribute and the shape is `[len(attrs)]`.
86+
context: Explicit context, if not from context manager.
87+
88+
Returns:
89+
DenseElementsAttr on success.
90+
91+
Raises:
92+
ValueError: If the type of the attributes does not match the type specified by `shaped_type`.
93+
)";
94+
7595
static const char kDenseResourceElementsAttrGetFromBufferDocstring[] =
7696
R"(Gets a DenseResourceElementsAttr from a Python buffer or array.
7797
@@ -647,6 +667,55 @@ class PyDenseElementsAttribute
647667
static constexpr const char *pyClassName = "DenseElementsAttr";
648668
using PyConcreteAttribute::PyConcreteAttribute;
649669

670+
static PyDenseElementsAttribute
671+
getFromList(py::list attributes, std::optional<PyType> explicitType,
672+
DefaultingPyMlirContext contextWrapper) {
673+
674+
const size_t numAttributes = py::len(attributes);
675+
if (numAttributes == 0)
676+
throw py::value_error("Attributes list must be non-empty");
677+
678+
MlirType shapedType;
679+
if (explicitType) {
680+
if ((!mlirTypeIsAShaped(*explicitType) ||
681+
!mlirShapedTypeHasStaticShape(*explicitType))) {
682+
std::string message =
683+
"Expected a static ShapedType for the shaped_type parameter: ";
684+
message.append(py::repr(py::cast(*explicitType)));
685+
throw py::value_error(message);
686+
}
687+
shapedType = *explicitType;
688+
} else {
689+
SmallVector<int64_t> shape{static_cast<int64_t>(numAttributes)};
690+
shapedType = mlirRankedTensorTypeGet(
691+
shape.size(), shape.data(),
692+
mlirAttributeGetType(pyTryCast<PyAttribute>(attributes[0])),
693+
mlirAttributeGetNull());
694+
}
695+
696+
SmallVector<MlirAttribute> mlirAttributes;
697+
mlirAttributes.reserve(numAttributes);
698+
for (auto attribute : attributes) {
699+
MlirAttribute mlirAttribute = pyTryCast<PyAttribute>(attribute);
700+
MlirType attrType = mlirAttributeGetType(mlirAttribute);
701+
mlirAttributes.push_back(mlirAttribute);
702+
703+
if (!mlirTypeEqual(mlirShapedTypeGetElementType(shapedType), attrType)) {
704+
std::string message = "All attributes must be of the same type and "
705+
"match the type parameter: expected=";
706+
message.append(py::repr(py::cast(shapedType)));
707+
message.append(", but got=");
708+
message.append(py::repr(py::cast(attrType)));
709+
throw py::value_error(message);
710+
}
711+
}
712+
713+
MlirAttribute elements = mlirDenseElementsAttrGet(
714+
shapedType, mlirAttributes.size(), mlirAttributes.data());
715+
716+
return PyDenseElementsAttribute(contextWrapper->getRef(), elements);
717+
}
718+
650719
static PyDenseElementsAttribute
651720
getFromBuffer(py::buffer array, bool signless,
652721
std::optional<PyType> explicitType,
@@ -883,6 +952,10 @@ class PyDenseElementsAttribute
883952
py::arg("type") = py::none(), py::arg("shape") = py::none(),
884953
py::arg("context") = py::none(),
885954
kDenseElementsAttrGetDocstring)
955+
.def_static("get", PyDenseElementsAttribute::getFromList,
956+
py::arg("attrs"), py::arg("type") = py::none(),
957+
py::arg("context") = py::none(),
958+
kDenseElementsAttrGetFromListDocstring)
886959
.def_static("get_splat", PyDenseElementsAttribute::getSplat,
887960
py::arg("shaped_type"), py::arg("element_attr"),
888961
"Gets a DenseElementsAttr where all values are the same")
@@ -954,8 +1027,8 @@ class PyDenseElementsAttribute
9541027
}
9551028
}; // namespace
9561029

957-
/// Refinement of the PyDenseElementsAttribute for attributes containing integer
958-
/// (and boolean) values. Supports element access.
1030+
/// Refinement of the PyDenseElementsAttribute for attributes containing
1031+
/// integer (and boolean) values. Supports element access.
9591032
class PyDenseIntElementsAttribute
9601033
: public PyConcreteAttribute<PyDenseIntElementsAttribute,
9611034
PyDenseElementsAttribute> {
@@ -964,8 +1037,8 @@ class PyDenseIntElementsAttribute
9641037
static constexpr const char *pyClassName = "DenseIntElementsAttr";
9651038
using PyConcreteAttribute::PyConcreteAttribute;
9661039

967-
/// Returns the element at the given linear position. Asserts if the index is
968-
/// out of range.
1040+
/// Returns the element at the given linear position. Asserts if the index
1041+
/// is out of range.
9691042
py::int_ dunderGetItem(intptr_t pos) {
9701043
if (pos < 0 || pos >= dunderLen()) {
9711044
throw py::index_error("attempt to access out of bounds element");
@@ -1267,7 +1340,8 @@ class PyStridedLayoutAttribute
12671340
return PyStridedLayoutAttribute(ctx->getRef(), attr);
12681341
},
12691342
py::arg("rank"), py::arg("context") = py::none(),
1270-
"Gets a strided layout attribute with dynamic offset and strides of a "
1343+
"Gets a strided layout attribute with dynamic offset and strides of "
1344+
"a "
12711345
"given rank.");
12721346
c.def_property_readonly(
12731347
"offset",

‎mlir/test/python/ir/array_attributes.py

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,87 @@ def testGetDenseElementsUnSupportedTypeOkIfExplicitTypeProvided():
5050
print(np.array(attr))
5151

5252

53+
################################################################################
54+
# Tests of the list of attributes .get() factory method
55+
################################################################################
56+
57+
58+
# CHECK-LABEL: TEST: testGetDenseElementsFromList
59+
@run
60+
def testGetDenseElementsFromList():
61+
with Context(), Location.unknown():
62+
attrs = [FloatAttr.get(F64Type.get(), 1.0), FloatAttr.get(F64Type.get(), 2.0)]
63+
attr = DenseElementsAttr.get(attrs)
64+
65+
# CHECK: dense<[1.000000e+00, 2.000000e+00]> : tensor<2xf64>
66+
print(attr)
67+
68+
69+
# CHECK-LABEL: TEST: testGetDenseElementsFromListWithExplicitType
70+
@run
71+
def testGetDenseElementsFromListWithExplicitType():
72+
with Context(), Location.unknown():
73+
attrs = [FloatAttr.get(F64Type.get(), 1.0), FloatAttr.get(F64Type.get(), 2.0)]
74+
shaped_type = ShapedType(Type.parse("tensor<2xf64>"))
75+
attr = DenseElementsAttr.get(attrs, shaped_type)
76+
77+
# CHECK: dense<[1.000000e+00, 2.000000e+00]> : tensor<2xf64>
78+
print(attr)
79+
80+
81+
# CHECK-LABEL: TEST: testGetDenseElementsFromListEmptyList
82+
@run
83+
def testGetDenseElementsFromListEmptyList():
84+
with Context(), Location.unknown():
85+
attrs = []
86+
87+
try:
88+
attr = DenseElementsAttr.get(attrs)
89+
except ValueError as e:
90+
# CHECK: Attributes list must be non-empty
91+
print(e)
92+
93+
94+
# CHECK-LABEL: TEST: testGetDenseElementsFromListNonAttributeType
95+
@run
96+
def testGetDenseElementsFromListNonAttributeType():
97+
with Context(), Location.unknown():
98+
attrs = [1.0]
99+
100+
try:
101+
attr = DenseElementsAttr.get(attrs)
102+
except RuntimeError as e:
103+
# CHECK: Invalid attribute when attempting to create an ArrayAttribute
104+
print(e)
105+
106+
107+
# CHECK-LABEL: TEST: testGetDenseElementsFromListMismatchedType
108+
@run
109+
def testGetDenseElementsFromListMismatchedType():
110+
with Context(), Location.unknown():
111+
attrs = [FloatAttr.get(F64Type.get(), 1.0), FloatAttr.get(F64Type.get(), 2.0)]
112+
shaped_type = ShapedType(Type.parse("tensor<2xf32>"))
113+
114+
try:
115+
attr = DenseElementsAttr.get(attrs, shaped_type)
116+
except ValueError as e:
117+
# CHECK: All attributes must be of the same type and match the type parameter
118+
print(e)
119+
120+
121+
# CHECK-LABEL: TEST: testGetDenseElementsFromListMixedTypes
122+
@run
123+
def testGetDenseElementsFromListMixedTypes():
124+
with Context(), Location.unknown():
125+
attrs = [FloatAttr.get(F64Type.get(), 1.0), FloatAttr.get(F32Type.get(), 2.0)]
126+
127+
try:
128+
attr = DenseElementsAttr.get(attrs)
129+
except ValueError as e:
130+
# CHECK: All attributes must be of the same type and match the type parameter
131+
print(e)
132+
133+
53134
################################################################################
54135
# Splats.
55136
################################################################################
@@ -205,6 +286,7 @@ def testGetDenseElementsBoolSplat():
205286

206287
### float and double arrays.
207288

289+
208290
# CHECK-LABEL: TEST: testGetDenseElementsF16
209291
@run
210292
def testGetDenseElementsF16():

0 commit comments

Comments
 (0)