1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
|
# RUN: %PYTHON %s | FileCheck %s
from mlir.ir import *
import mlir.dialects.arith as arith
import mlir.dialects.builtin as builtin
import mlir.dialects.tensor as tensor
def run(f):
print("\nTEST:", f.__name__)
f()
return f
# CHECK-LABEL: TEST: testDimOp
@run
def testDimOp():
with Context() as ctx, Location.unknown():
module = Module.create()
f32Type = F32Type.get()
indexType = IndexType.get()
with InsertionPoint(module.body):
@builtin.FuncOp.from_py_func(RankedTensorType.get((-1, -1), f32Type))
# CHECK: func @tensor_static_dim
# CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>
# CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
# CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
# CHECK: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
# CHECK: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
# CHECK: return %[[D0]], %[[D1]]
def tensor_static_dim(t):
c0 = arith.ConstantOp(indexType, 0)
c1 = arith.ConstantOp(indexType, 1)
d0 = tensor.DimOp(t, c0)
d1 = tensor.DimOp(t, c1)
return [d0.result, d1.result]
print(module)
|
