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