1// RUN: mlir-linalg-ods-gen %s -gen-ods-decl=1 | FileCheck %s --check-prefix=ODS
2// RUN: mlir-linalg-ods-gen %s -gen-impl=1 | FileCheck %s --check-prefix=IMPL
3
4// ODS-LABEL: def Test1Op : LinalgStructuredBase_Op<"test1", [
5//  ODS-NEXT:   AttrSizedOperandSegments
6//  ODS-NEXT:   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
7//  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
8//
9// IMPL-LABEL:  ArrayAttr Test1Op::iterator_types() {
10//       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
11//
12//       IMPL:  ArrayAttr Test1Op::indexing_maps() {
13//       IMPL: auto s0 = getAffineSymbolExpr(0, context); (void)s0;
14//  IMPL-NEXT: auto s1 = getAffineSymbolExpr(1, context); (void)s1;
15//  IMPL-NEXT: auto map0 = AffineMap::get(2, 2, {d0, d1}, context);
16//  IMPL-NEXT: map0 = map0.replaceDimsAndSymbols({}, { s0, s1 }, 2, 0);
17//  IMPL-NEXT: map0 = simplifyAffineMap(map0);
18//  IMPL-NEXT: auto map1 = AffineMap::get(2, 2, {d1}, context);
19//  IMPL-NEXT: map1 = map1.replaceDimsAndSymbols({}, { s0, s1 }, 2, 0);
20//  IMPL-NEXT: map1 = simplifyAffineMap(map1);
21//  IMPL-NEXT: auto map2 = AffineMap::get(2, 2, {d0}, context);
22//  IMPL-NEXT: map2 = map2.replaceDimsAndSymbols({}, { s0, s1 }, 2, 0);
23//  IMPL-NEXT: map2 = simplifyAffineMap(map2);
24//  IMPL-NEXT: return {{.+}}.getAffineMapArrayAttr({ map0, map1, map2 });
25//
26//       IMPL:  void Test1Op::regionBuilder(Block &block) {
27//       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
28//       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
29//       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
30//       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
31//
32ods_def<Test1Op> :
33def test1(A: f32(M, K), B: f32(K)) -> (C: f32(M)) {
34  C(m) = std_addf<k>(std_mulf(A(m, k), B(k)));
35}
36
37// ODS-LABEL: def Test2Op : LinalgStructuredBase_Op<"test2", [
38//  ODS-NEXT:   AttrSizedOperandSegments
39//  ODS-NEXT:   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
40//  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
41//
42// IMPL-LABEL:  ArrayAttr Test2Op::iterator_types() {
43//       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
44//
45//       IMPL:  ArrayAttr Test2Op::indexing_maps() {
46//       IMPL:  AffineMap::get(3, 3, {d0, d2}, context)
47//       IMPL:  AffineMap::get(3, 3, {d2, d1}, context)
48//       IMPL:  AffineMap::get(3, 3, {d0, d1}, context)
49//
50//       IMPL:  Test2Op::regionBuilder(Block &block) {
51//       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
52//       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
53//       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
54//       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
55//
56ods_def<Test2Op> :
57def test2(A: f32(M, K), B: f32(K, N)) -> (C: f32(M, N)) {
58  C(m, n) = std_addf<k>(std_mulf(A(m, k), B(k, n)));
59}
60
61// ODS-LABEL: def Test3Op : LinalgStructuredBase_Op<"test3", [
62//  ODS-NEXT:   AttrSizedOperandSegments
63//  ODS-NEXT:   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
64//  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
65//
66// IMPL-LABEL:  ArrayAttr Test3Op::iterator_types() {
67//       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
68//
69//       IMPL:  ArrayAttr Test3Op::indexing_maps() {
70//       IMPL:  AffineMap::get(4, 4, {d0, d1, d3}, context)
71//       IMPL:  AffineMap::get(4, 4, {d3, d2}, context)
72//       IMPL:  AffineMap::get(4, 4, {d0, d1, d2}, context)
73//
74//       IMPL:  Test3Op::regionBuilder(Block &block) {
75//       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
76//       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
77//       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
78//       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
79//
80ods_def<Test3Op> :
81def test3(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N)) {
82  C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
83}
84
85// Test attribute definitions
86// ODS-LABEL: def Test4Op
87// ODS: F32ArrayAttr:$array_attr,
88// ODS: F32:$f32_attr,
89// ODS: RankedF32ElementsAttr<[4]>:$fvec_attr,
90// ODS: I32:$i32_attr,
91// ODS: I64:$i64_attr,
92// ODS: RankedI32ElementsAttr<[5, 6]>:$ivec_attr,
93// ODS: OptionalAttr<F32>:$optional_attr
94//
95ods_def<Test4Op> :
96def test4(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N))
97attr(
98  f32_attr: f32,
99  i32_attr: i32,
100  i64_attr: i64,
101  fvec_attr: 4xf32,
102  ivec_attr: 5x6xi32,
103  array_attr : f32[],
104  optional_attr? : f32
105) {
106  C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
107}
108
109// Test attribute usage in affine expressions
110// IMPL-LABEL: ArrayAttr Test5Op::indexing_maps() {
111// IMPL: auto cst0 = getAffineConstantExpr(strides().getValue<int>({ 0 }), context);
112// IMPL: auto cst1 = getAffineConstantExpr(strides().getValue<int>({ 1 }), context);
113// IMPL: auto map0 = AffineMap::get(7, 9, {d0, d1 * s7 + d4, d2 * s8 + d5, d6}, context);
114// IMPL: map0 = map0.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);
115// IMPL: map0 = simplifyAffineMap(map0);
116// IMPL: auto map1 = AffineMap::get(7, 9, {d3, d4, d5, d6}, context);
117// IMPL: map1 = map1.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);
118// IMPL: map1 = simplifyAffineMap(map1);
119// IMPL: auto map2 = AffineMap::get(7, 7, {d0, d1, d2, d3}, context);
120// IMPL: map2 = map2.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);
121// IMPL: map2 = simplifyAffineMap(map2);
122// IMPL: return {{.+}}.getAffineMapArrayAttr({ map0, map1, map2 });
123//
124ods_def<Test5Op>:
125def test5(I: f32(N, H, W, C), K: f32(F, KH, KW, C)) -> (O: f32(N, H, W, F))
126     attr(strides: 2xi32) {
127  O(n, h, w, f) = std_addf<kh, kw>(std_mulf(
128    I(n, h * strides[0] + kh, w * strides[1] + kw, c), K(f, kh, kw, c)));
129}
130
131// Test documentation
132// ODS-LABEL: def Test6Op
133// ODS:       let summary = [{ My magic op. }];
134// ODS-NEXT:  let description = [{
135// ODS-NEXT:    It has two inputs.
136// ODS-NEXT:    It has one output.
137// ODS-NEXT:  }];
138//
139ods_def<Test6Op>:
140def test6(A: f32(M, K), B: f32(K)) -> (C: f32(M))
141"""
142My magic op.
143
144It has two inputs.
145It has one output.
146"""
147{
148  C(m) = std_addf<k>(std_mulf(A(m, k), B(k)));
149}
150
151// Test attribute builder
152// ODS-LABEL: def Test7Op
153// ODS:         OpBuilderDAG<
154// ODS:           (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
155// ODS:            "ValueRange":$outputs, "Attribute":$attr_a, "Attribute":$attr_b)
156// ODS:           $_state.addAttribute("attr_a", attr_a);
157// ODS:           $_state.addAttribute("attr_b", attr_b);
158//
159ods_def<Test7Op>:
160def test7(A: f32(M, K), B: f32(K)) -> (C: f32(M))
161     attr(attr_a: f32, attr_b: 4xi32)
162{
163  C(m) = std_addf<k>(std_mulf(A(m, k), B(k)));
164}
165