llvm-capstone/mlir/test/python/dialects/shape.py
Tobias Hieta f9008e6366
[NFC][Py Reformat] Reformat python files in mlir subdir
This is an ongoing series of commits that are reformatting our
Python code.

Reformatting is done with `black`.

If you end up having problems merging this commit because you
have made changes to a python file, the best way to handle that
is to run git checkout --ours <yourfile> and then reformat it
with black.

If you run into any problems, post to discourse about it and
we will try to help.

RFC Thread below:

https://discourse.llvm.org/t/rfc-document-and-standardize-python-code-style

Differential Revision: https://reviews.llvm.org/D150782
2023-05-26 08:05:40 +02:00

45 lines
1.4 KiB
Python

# RUN: %PYTHON %s | FileCheck %s
from mlir.ir import *
import numpy as np
import mlir.dialects.func as func
import mlir.dialects.shape as shape
def run(f):
print("\nTEST:", f.__name__)
f()
return f
# CHECK-LABEL: TEST: testConstShape
@run
def testConstShape():
with Context() as ctx, Location.unknown():
module = Module.create()
f32 = F32Type.get()
with InsertionPoint(module.body):
@func.FuncOp.from_py_func(
RankedTensorType.get((12, ShapedType.get_dynamic_size()), f32)
)
def const_shape_tensor(arg):
shape.ConstWitnessOp(False)
shape.ConstSizeOp(30)
shape.ConstSizeOp(IntegerAttr.get(IndexType.get(), 40))
x = shape.ConstShapeOp([1, 2])
shape.MeetOp(x, x, error="impossible")
return shape.ConstShapeOp(
DenseElementsAttr.get(
np.array([3, 4], dtype=np.int64), type=IndexType.get()
)
)
# CHECK-LABEL: func @const_shape_tensor(%arg0: tensor<12x?xf32>)
# CHECK-DAG: shape.const_witness false
# CHECK-DAG: shape.const_size 30
# CHECK-DAG: shape.const_size 40
# CHECK-DAG: shape.const_shape [1, 2] : tensor<2xindex>
# CHECK-DAG: shape.const_shape [3, 4] : tensor<2xindex>
print(module)