Perry Gibson d898ff650a
[mlir,python] Fix case when FuncOp.arg_attrs is not set (#117188)
FuncOps can have `arg_attrs`, an array of dictionary attributes
associated with their arguments.

E.g., 

```mlir
func.func @main(%arg0: tensor<8xf32> {test.attr_name = "value"}, %arg1: tensor<8x16xf32>)
```

These are exposed via the MLIR Python bindings with
`my_funcop.arg_attrs`.

In this case, it would return `[{test.attr_name = "value"}, {}]`, i.e.,
`%arg1` has an empty `DictAttr`.

However, if I try and access this property from a FuncOp with an empty
`arg_attrs`, e.g.,

```mlir
func.func @main(%arg0: tensor<8xf32>, %arg1: tensor<8x16xf32>)
```

This raises the error:

```python
    return ArrayAttr(self.attributes[ARGUMENT_ATTRIBUTE_NAME])
                     ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'attempt to access a non-existent attribute'
```

This PR fixes this by returning the expected `[{}, {}]`.
2024-12-02 08:55:51 -08:00

136 lines
3.8 KiB
Python

# RUN: %PYTHON %s | FileCheck %s
from mlir.ir import *
from mlir.dialects import arith
from mlir.dialects import builtin
from mlir.dialects import func
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: TEST: testConstantOp
@constructAndPrintInModule
def testConstantOp():
c1 = arith.ConstantOp(IntegerType.get_signless(32), 42)
c2 = arith.ConstantOp(IntegerType.get_signless(64), 100)
c3 = arith.ConstantOp(F32Type.get(), 3.14)
c4 = arith.ConstantOp(F64Type.get(), 1.23)
# CHECK: 42
print(c1.literal_value)
# CHECK: 100
print(c2.literal_value)
# CHECK: 3.140000104904175
print(c3.literal_value)
# CHECK: 1.23
print(c4.literal_value)
# CHECK: = arith.constant 42 : i32
# CHECK: = arith.constant 100 : i64
# CHECK: = arith.constant 3.140000e+00 : f32
# CHECK: = arith.constant 1.230000e+00 : f64
# CHECK-LABEL: TEST: testVectorConstantOp
@constructAndPrintInModule
def testVectorConstantOp():
int_type = IntegerType.get_signless(32)
vec_type = VectorType.get([2, 2], int_type)
c1 = arith.ConstantOp(
vec_type, DenseElementsAttr.get_splat(vec_type, IntegerAttr.get(int_type, 42))
)
try:
print(c1.literal_value)
except ValueError as e:
assert "only integer and float constants have literal values" in str(e)
else:
assert False
# CHECK: = arith.constant dense<42> : vector<2x2xi32>
# CHECK-LABEL: TEST: testConstantIndexOp
@constructAndPrintInModule
def testConstantIndexOp():
c1 = arith.ConstantOp.create_index(10)
# CHECK: 10
print(c1.literal_value)
# CHECK: = arith.constant 10 : index
# CHECK-LABEL: TEST: testFunctionCalls
@constructAndPrintInModule
def testFunctionCalls():
foo = func.FuncOp("foo", ([], []))
foo.sym_visibility = StringAttr.get("private")
bar = func.FuncOp("bar", ([], [IndexType.get()]))
bar.sym_visibility = StringAttr.get("private")
qux = func.FuncOp("qux", ([], [F32Type.get()]))
qux.sym_visibility = StringAttr.get("private")
con = func.ConstantOp(qux.type, qux.sym_name.value)
assert con.type == qux.type
with InsertionPoint(func.FuncOp("caller", ([], [])).add_entry_block()):
func.CallOp(foo, [])
func.CallOp([IndexType.get()], "bar", [])
func.CallOp([F32Type.get()], FlatSymbolRefAttr.get("qux"), [])
func.ReturnOp([])
# CHECK: func private @foo()
# CHECK: func private @bar() -> index
# CHECK: func private @qux() -> f32
# CHECK: %f = func.constant @qux : () -> f32
# CHECK: func @caller() {
# CHECK: call @foo() : () -> ()
# CHECK: %0 = call @bar() : () -> index
# CHECK: %1 = call @qux() : () -> f32
# CHECK: return
# CHECK: }
# CHECK-LABEL: TEST: testFunctionArgAttrs
@constructAndPrintInModule
def testFunctionArgAttrs():
foo = func.FuncOp("foo", ([F32Type.get()], []))
foo.sym_visibility = StringAttr.get("private")
foo2 = func.FuncOp("foo2", ([F32Type.get(), F32Type.get()], []))
foo2.sym_visibility = StringAttr.get("private")
empty_attr = DictAttr.get({})
test_attr = DictAttr.get({"test.foo": StringAttr.get("bar")})
test_attr2 = DictAttr.get({"test.baz": StringAttr.get("qux")})
assert len(foo.arg_attrs) == 1
assert foo.arg_attrs[0] == empty_attr
foo.arg_attrs = [test_attr]
assert foo.arg_attrs[0]["test.foo"] == StringAttr.get("bar")
assert len(foo2.arg_attrs) == 2
assert foo2.arg_attrs == ArrayAttr.get([empty_attr, empty_attr])
foo2.arg_attrs = [empty_attr, test_attr2]
assert foo2.arg_attrs == ArrayAttr.get([empty_attr, test_attr2])
# CHECK: func private @foo(f32 {test.foo = "bar"})
# CHECK: func private @foo2(f32, f32 {test.baz = "qux"})