Fix array attribute in bindings for linalg.init_tensor

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D101998
This commit is contained in:
Denys Shabalin
2021-05-06 18:24:07 +02:00
committed by Alex Zinenko
parent 1c4cb510b4
commit 1f109f9d9c
2 changed files with 14 additions and 3 deletions

View File

@@ -74,9 +74,9 @@ class InitTensorOp:
result_type = RankedTensorType.get(sizes, element_type)
static_size_ints = sizes
index_type = IndexType.get(context)
i64_type = IntegerType.get_signless(64)
attributes["static_sizes"] = ArrayAttr.get(
[IntegerAttr.get(index_type, s) for s in static_size_ints],
[IntegerAttr.get(i64_type, s) for s in static_size_ints],
context=context)
op = self.build_generic(results=[result_type],
operands=operands,

View File

@@ -38,6 +38,17 @@ def testInitTensor():
print(module)
# CHECK-LABEL: TEST: testInitTensorStaticSizesAttribute
@run
def testInitTensorStaticSizesAttribute():
with Context() as ctx, Location.unknown():
module = Module.create()
f32 = F32Type.get()
with InsertionPoint(module.body):
op = linalg.InitTensorOp([3, 4], f32)
# CHECK: [3, 4]
print(op.attributes['static_sizes'])
# CHECK-LABEL: TEST: testFill
@run
def testFill():
@@ -153,7 +164,7 @@ def testNamedStructuredOpGenericForm():
# CHECK-NEXT: std.mulf{{.*}} (f32, f32) -> f32
# CHECK-NEXT: std.addf{{.*}} (f32, f32) -> f32
# CHECK-NEXT: linalg.yield{{.*}} (f32) -> ()
# CHECK-NEXT: {linalg.memoized_indexing_maps{{.*}}operand_segment_sizes = dense<[2, 1]> : vector<2xi32>} :
# CHECK-NEXT: {linalg.memoized_indexing_maps{{.*}}operand_segment_sizes = dense<[2, 1]> : vector<2xi32>} :
# CHECK-SAME: (tensor<4x16xf32>, tensor<16x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32>
return linalg.matmul(lhs, rhs, outs=[init_result.result])