1// RUN: mlir-opt %s -split-input-file -linalg-fold-unit-extent-dims="fold-one-trip-loops-only" | FileCheck %s
2
3#accesses = [
4  affine_map<(i, j, k, l, m) -> (i, k, m)>,
5  affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>
6]
7
8#trait = {
9  iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],
10  indexing_maps = #accesses,
11  library_call = "some_external_func"
12}
13
14func @drop_one_trip_loops(%arg0 : tensor<?x1x?xf32>, %shape: tensor<?x1x?x1x?xf32>) -> tensor<?x1x?x1x?xf32>
15{
16  %0 = linalg.generic #trait
17    ins(%arg0 : tensor<?x1x?xf32>)
18    outs(%shape : tensor<?x1x?x1x?xf32>) {
19       ^bb0(%arg1 : f32, %arg2 : f32) :
20         linalg.yield %arg1 : f32
21       } -> tensor<?x1x?x1x?xf32>
22  return %0 : tensor<?x1x?x1x?xf32>
23}
24//   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d2)>
25//   CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d1, 0, d2)>
26// CHECK-LABEL: func @drop_one_trip_loops
27//       CHECK:   linalg.generic
28//  CHECK-SAME:     indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
29//  CHECK-SAME:     iterator_types = ["parallel", "parallel", "parallel"]
30
31// -----
32
33#map0 = affine_map<(i, j) -> (i, j)>
34#access = [#map0, #map0]
35#trait = {
36  iterator_types = ["parallel", "parallel"],
37  indexing_maps = #access,
38  library_call = "some_external_func"
39}
40
41func @drop_all_loops(%arg0 : tensor<1x1xf32>) -> tensor<1x1xf32>
42{
43  %0 = linalg.generic #trait
44     ins(%arg0 : tensor<1x1xf32>)
45    outs(%arg0 : tensor<1x1xf32>) {
46       ^bb0(%arg1: f32, %arg2: f32) :
47         linalg.yield %arg1 : f32
48       } -> tensor<1x1xf32>
49  return %0 : tensor<1x1xf32>
50}
51//   CHECK-DAG: #[[$MAP0:.*]] = affine_map<() -> (0, 0)>
52// CHECK-LABEL: func @drop_all_loops
53//       CHECK:   linalg.generic
54//  CHECK-SAME:     indexing_maps = [#[[$MAP0]], #[[$MAP0]]]
55//  CHECK-SAME:     iterator_types = []
56
57// -----
58
59#map0 = affine_map<(i, j) -> (i, j)>
60#access = [#map0, #map0]
61#trait = {
62  iterator_types = ["parallel", "parallel"],
63  indexing_maps = #access,
64  library_call = "some_external_func"
65}
66
67func @drop_all_loops(%arg0 : memref<1x1xf32>, %arg1 : memref<1x1xf32>)
68{
69  linalg.generic #trait
70     ins(%arg0 : memref<1x1xf32>)
71    outs(%arg1 : memref<1x1xf32>) {
72    ^bb0(%arg2: f32, %arg3 : f32) :
73      linalg.yield %arg2 : f32
74    }
75  return
76}
77//   CHECK-DAG: #[[$MAP0:.*]] = affine_map<() -> (0, 0)>
78// CHECK-LABEL: func @drop_all_loops
79//       CHECK:   linalg.generic
80//  CHECK-SAME:     indexing_maps = [#[[$MAP0]], #[[$MAP0]]]
81//  CHECK-SAME:     iterator_types = []
82
83// -----
84
85#accesses = [
86  affine_map<(d0, d1) -> (d0, d1)>,
87  affine_map<(d0, d1) -> (d1)>
88]
89
90#trait = {
91  indexing_maps = #accesses,
92  iterator_types = ["parallel", "parallel"],
93  library_call = "some_external_fn"
94}
95
96func @leading_dim_1_canonicalization(%arg0: tensor<1x5xf32>, %shape: tensor<5xf32>) -> tensor<5xf32> {
97  %0 = linalg.generic #trait
98       ins(%arg0 : tensor<1x5xf32>)
99      outs(%shape : tensor<5xf32>) {
100    ^bb0(%arg2: f32, %arg3: f32):     // no predecessors
101      linalg.yield %arg2 : f32
102  } -> tensor<5xf32>
103  return %0 : tensor<5xf32>
104}
105//   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (0, d0)>
106//   CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
107// CHECK-LABEL: func @leading_dim_1_canonicalization
108//       CHECK:   linalg.generic
109//  CHECK-SAME:     indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
110//  CHECK-SAME:     iterator_types = ["parallel"]
111