diff --git a/mlir/test/Dialect/Linalg/roundtrip.mlir b/mlir/test/Dialect/Linalg/roundtrip.mlir --- a/mlir/test/Dialect/Linalg/roundtrip.mlir +++ b/mlir/test/Dialect/Linalg/roundtrip.mlir @@ -310,6 +310,41 @@ // ----- +func @foo(%i: index, %j: index, %k: index, %0: vector<3x4xi4>, %1: f32) -> f32 { + %f0 = constant 0.0 : f32 + return %f0 : f32 +} + +#accesses = [ + affine_map<(i, j, k) -> (j, i)>, + affine_map<(i, j, k) -> (i, k, i + j)> +] + +#trait2 = { + args_in = 2, + args_out = 1, + indexing_maps = #accesses, + iterator_types = ["parallel", "parallel", "parallel"], + fun = @foo, + library_call = "some_external_function_name_1" +} + +func @indexed_generic_with_tensor_input_and_output( + %arg0: tensor>, %arg1: tensor) + -> (tensor) { + %0 = linalg.indexed_generic #trait2 %arg0, %arg1 {foo = 1} : + tensor>, tensor -> tensor + return %0 : tensor +} +// CHECK-LABEL: func @indexed_generic_with_tensor_input_and_output +// CHECK: linalg.indexed_generic {args_in = 2 : i64, args_out = 1 : i64, fun = @foo, +// CHECK-SAME: indexing_maps = [#{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel"], +// CHECK-SAME: library_call = "some_external_function_name_1"} %{{.*}}, %{{.*}} {foo = 1 : i64}: +// CHECK-SAME: tensor>, tensor -> tensor +// CHECK: return {{.*}} : tensor + +// ----- + // CHECK-DAG: #[[strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)> // CHECK-DAG: #[[strided3D:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>