xref: /llvm-project/mlir/test/Dialect/Linalg/tile-indexed.mlir (revision d28a4f1fc02dc34a87fa22af0a053e8f1e7f6cea)
1// RUN: mlir-opt %s -transform-interpreter -canonicalize -split-input-file | FileCheck %s -check-prefix=TILE-10n25
2
3func.func @indexed_vector(%arg0: memref<50xindex>) {
4  linalg.generic {indexing_maps = [affine_map<(i) -> (i)>],
5                  iterator_types = ["parallel"]}
6     outs(%arg0 : memref<50xindex>) {
7    ^bb0(%a: index):
8      %i = linalg.index 0 : index
9      linalg.yield %i : index
10  }
11  return
12}
13
14module attributes {transform.with_named_sequence} {
15  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
16      %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
17      %1, %loop = transform.structured.tile_using_for %0 tile_sizes [10] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
18      transform.yield
19  }
20}
21
22// TILE-10n25-DAG: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0)[s0] -> (d0 + s0)>
23// TILE-10n25-LABEL: func @indexed_vector
24// TILE-10n25: %[[C10:.*]] = arith.constant 10 : index
25// TILE-10n25: scf.for %[[J:.*]] = {{.*}} step %[[C10]]
26// TILE-10n25:   linalg.generic
27// TILE-10n25:     %[[I:.*]] = linalg.index 0 : index
28// TILE-10n25:     %[[NEW_I:.*]] = affine.apply [[$MAP]](%[[J]])[%[[I]]]
29// TILE-10n25:     linalg.yield %[[NEW_I]] : index
30
31// -----
32
33func.func @indexed_matrix(%arg0: memref<50x50xindex>) {
34  linalg.generic {indexing_maps = [affine_map<(i, j) -> (i, j)>],
35                  iterator_types = ["parallel", "parallel"]}
36    outs(%arg0 : memref<50x50xindex>) {
37    ^bb0(%a: index):
38      %i = linalg.index 0 : index
39      %j = linalg.index 1 : index
40      %sum = arith.addi %i, %j : index
41      linalg.yield %sum : index
42  }
43  return
44}
45
46module attributes {transform.with_named_sequence} {
47  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
48      %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
49      %1, %loop:2 = transform.structured.tile_using_for %0 tile_sizes [10, 25] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
50      transform.yield
51  }
52}
53
54// TILE-10n25-DAG: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0)[s0] -> (d0 + s0)>
55// TILE-10n25-LABEL: func @indexed_matrix
56// TILE-10n25-DAG: %[[C25:.*]] = arith.constant 25 : index
57// TILE-10n25-DAG: %[[C10:.*]] = arith.constant 10 : index
58// TILE-10n25: scf.for %[[K:.*]] = {{.*}} step %[[C10]]
59// TILE-10n25:   scf.for %[[L:.*]] = {{.*}} step %[[C25]]
60// TILE-10n25:     linalg.generic
61// TILE-10n25:       %[[I:.*]] = linalg.index 0 : index
62// TILE-10n25:       %[[NEW_I:.*]] = affine.apply [[$MAP]](%[[K]])[%[[I]]]
63// TILE-10n25:       %[[J:.*]] = linalg.index 1 : index
64// TILE-10n25:       %[[NEW_J:.*]] = affine.apply [[$MAP]](%[[L]])[%[[J]]]
65// TILE-10n25:       %[[SUM:.*]] = arith.addi %[[NEW_I]], %[[NEW_J]] : index
66// TILE-10n25:       linalg.yield %[[SUM]] : index
67