forked from OSchip/llvm-project
210 lines
8.4 KiB
Plaintext
210 lines
8.4 KiB
Plaintext
// RUN: mlir-linalg-ods-gen %s -gen-ods-decl=1 | FileCheck %s --check-prefix=ODS
|
|
// RUN: mlir-linalg-ods-gen %s -gen-impl=1 | FileCheck %s --check-prefix=IMPL
|
|
|
|
// ODS-LABEL: def Test1Op : LinalgStructuredBase_Op<"test1", [
|
|
// ODS-NEXT: AttrSizedOperandSegments
|
|
// ODS-NEXT: DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
|
|
// ODS-NEXT: SingleBlockImplicitTerminator<"YieldOp">
|
|
//
|
|
// IMPL-LABEL: ArrayAttr Test1Op::iterator_types() {
|
|
// IMPL: { {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
|
|
//
|
|
// IMPL: ArrayAttr Test1Op::indexing_maps() {
|
|
// IMPL: auto s0 = getAffineSymbolExpr(0, context); (void)s0;
|
|
// IMPL-NEXT: auto s1 = getAffineSymbolExpr(1, context); (void)s1;
|
|
// IMPL-NEXT: auto map0 = AffineMap::get(2, 2, {d0, d1}, context);
|
|
// IMPL-NEXT: map0 = map0.replaceDimsAndSymbols({}, { s0, s1 }, 2, 0);
|
|
// IMPL-NEXT: map0 = simplifyAffineMap(map0);
|
|
// IMPL-NEXT: auto map1 = AffineMap::get(2, 2, {d1}, context);
|
|
// IMPL-NEXT: map1 = map1.replaceDimsAndSymbols({}, { s0, s1 }, 2, 0);
|
|
// IMPL-NEXT: map1 = simplifyAffineMap(map1);
|
|
// IMPL-NEXT: auto map2 = AffineMap::get(2, 2, {d0}, context);
|
|
// IMPL-NEXT: map2 = map2.replaceDimsAndSymbols({}, { s0, s1 }, 2, 0);
|
|
// IMPL-NEXT: map2 = simplifyAffineMap(map2);
|
|
// IMPL-NEXT: return {{.+}}.getAffineMapArrayAttr({ map0, map1, map2 });
|
|
//
|
|
// IMPL: void Test1Op::regionBuilder(ImplicitLocOpBuilder &b,
|
|
// IMPL: Block &block) {
|
|
// IMPL: Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
|
|
// IMPL: Value [[d:.*]] = b.create<MulFOp>([[a]], [[b]]);
|
|
// IMPL: Value [[e:.*]] = b.create<AddFOp>([[c]], [[d]]);
|
|
// IMPL: b.create<linalg::YieldOp>(ValueRange{ [[e]] });
|
|
//
|
|
ods_def<Test1Op> :
|
|
def test1(A: f32(M, K), B: f32(K)) -> (C: f32(M)) {
|
|
C(m) = AddFOp<k>(C(m), MulFOp(A(m, k), B(k)));
|
|
}
|
|
|
|
// ODS-LABEL: def Test2Op : LinalgStructuredBase_Op<"test2", [
|
|
// ODS-NEXT: AttrSizedOperandSegments
|
|
// ODS-NEXT: DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
|
|
// ODS-NEXT: SingleBlockImplicitTerminator<"YieldOp">
|
|
//
|
|
// IMPL-LABEL: ArrayAttr Test2Op::iterator_types() {
|
|
// IMPL: { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
|
|
//
|
|
// IMPL: ArrayAttr Test2Op::indexing_maps() {
|
|
// IMPL: AffineMap::get(3, 3, {d0, d2}, context)
|
|
// IMPL: AffineMap::get(3, 3, {d2, d1}, context)
|
|
// IMPL: AffineMap::get(3, 3, {d0, d1}, context)
|
|
//
|
|
// IMPL: Test2Op::regionBuilder(ImplicitLocOpBuilder &b,
|
|
// IMPL: Block &block) {
|
|
// IMPL: Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
|
|
// IMPL: Value [[d:.*]] = b.create<MulFOp>([[a]], [[b]]);
|
|
// IMPL: Value [[e:.*]] = b.create<AddFOp>([[c]], [[d]]);
|
|
// IMPL: b.create<linalg::YieldOp>(ValueRange{ [[e]] });
|
|
//
|
|
ods_def<Test2Op> :
|
|
def test2(A: f32(M, K), B: f32(K, N)) -> (C: f32(M, N)) {
|
|
C(m, n) = AddFOp<k>(C(m, n), MulFOp(A(m, k), B(k, n)));
|
|
}
|
|
|
|
// ODS-LABEL: def Test3Op : LinalgStructuredBase_Op<"test3", [
|
|
// ODS-NEXT: AttrSizedOperandSegments
|
|
// ODS-NEXT: DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
|
|
// ODS-NEXT: SingleBlockImplicitTerminator<"YieldOp">
|
|
//
|
|
// IMPL-LABEL: ArrayAttr Test3Op::iterator_types() {
|
|
// IMPL: { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
|
|
//
|
|
// IMPL: ArrayAttr Test3Op::indexing_maps() {
|
|
// IMPL: AffineMap::get(4, 4, {d0, d1, d3}, context)
|
|
// IMPL: AffineMap::get(4, 4, {d3, d2}, context)
|
|
// IMPL: AffineMap::get(4, 4, {d0, d1, d2}, context)
|
|
//
|
|
// IMPL: Test3Op::regionBuilder(ImplicitLocOpBuilder &b,
|
|
// IMPL: Block &block) {
|
|
// IMPL: Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
|
|
// IMPL: Value [[d:.*]] = b.create<MulFOp>([[a]], [[b]]);
|
|
// IMPL: Value [[e:.*]] = b.create<AddFOp>([[c]], [[d]]);
|
|
// IMPL: b.create<linalg::YieldOp>(ValueRange{ [[e]] });
|
|
//
|
|
ods_def<Test3Op> :
|
|
def test3(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N)) {
|
|
C(b, m, n) = AddFOp<k>(C(b, m, n), MulFOp(A(b, m, k), B(k, n)));
|
|
}
|
|
|
|
// Test attribute definitions
|
|
// ODS-LABEL: def Test4Op
|
|
// ODS: F32ArrayAttr:$array_attr,
|
|
// ODS: F32Attr:$f32_attr,
|
|
// ODS: RankedF32ElementsAttr<[4]>:$fvec_attr,
|
|
// ODS: I32Attr:$i32_attr,
|
|
// ODS: I64Attr:$i64_attr,
|
|
// ODS: RankedI32ElementsAttr<[5, 6]>:$ivec_attr,
|
|
// ODS: OptionalAttr<F32Attr>:$optional_attr
|
|
//
|
|
// ODS: bool hasDynamicIndexingMaps();
|
|
// ODS: LogicalResult verifyIndexingMapRequiredAttributes();
|
|
//
|
|
// IMPL: bool Test4Op::hasDynamicIndexingMaps() { return true; }
|
|
// IMPL: LogicalResult Test4Op::verifyIndexingMapRequiredAttributes()
|
|
// IMPL: op->getAttrOfType<ArrayAttr>("array_attr")
|
|
// IMPL: op->getAttr("f32_attr")
|
|
// IMPL: op->getAttrOfType<DenseElementsAttr>("fvec_attr")
|
|
// IMPL: op->getAttr("i32_attr")
|
|
// IMPL: op->getAttr("i64_attr")
|
|
// IMPL: op->getAttrOfType<DenseElementsAttr>("ivec_attr")
|
|
//
|
|
ods_def<Test4Op> :
|
|
def test4(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N))
|
|
attr(
|
|
f32_attr: f32,
|
|
i32_attr: i32,
|
|
i64_attr: i64,
|
|
fvec_attr: 4xf32,
|
|
ivec_attr: 5x6xi32,
|
|
array_attr : f32[],
|
|
optional_attr? : f32
|
|
) {
|
|
C(b, m, n) = AddFOp<k>(C(b, m, n), MulFOp(A(b, m, k), B(k, n)));
|
|
}
|
|
|
|
// Test attribute usage in affine expressions
|
|
// IMPL-LABEL: ArrayAttr Test5Op::indexing_maps() {
|
|
// IMPL: auto cst0 = getAffineConstantExpr(strides().getValue<int>({ 0 }), context);
|
|
// IMPL: auto cst1 = getAffineConstantExpr(strides().getValue<int>({ 1 }), context);
|
|
// IMPL: auto map0 = AffineMap::get(7, 9, {d0, d1 * s7 + d4, d2 * s8 + d5, d6}, context);
|
|
// IMPL: map0 = map0.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);
|
|
// IMPL: map0 = simplifyAffineMap(map0);
|
|
// IMPL: auto map1 = AffineMap::get(7, 9, {d3, d4, d5, d6}, context);
|
|
// IMPL: map1 = map1.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);
|
|
// IMPL: map1 = simplifyAffineMap(map1);
|
|
// IMPL: auto map2 = AffineMap::get(7, 7, {d0, d1, d2, d3}, context);
|
|
// IMPL: map2 = map2.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);
|
|
// IMPL: map2 = simplifyAffineMap(map2);
|
|
// IMPL: return {{.+}}.getAffineMapArrayAttr({ map0, map1, map2 });
|
|
//
|
|
ods_def<Test5Op>:
|
|
def test5(I: f32(N, H, W, C), K: f32(F, KH, KW, C)) -> (O: f32(N, H, W, F))
|
|
attr(strides: 2xi32) {
|
|
O(n, h, w, f) = AddFOp<kh, kw>(
|
|
MulFOp(AddFOp(I(n, h * strides[0] + kh, w * strides[1] + kw, c),
|
|
I(n, h * strides[0] + kh, w * strides[1] + kw, c)),
|
|
K(f, kh, kw, c)));
|
|
}
|
|
|
|
// Test documentation
|
|
// ODS-LABEL: def Test6Op
|
|
// ODS: let summary = [{ My magic op. }];
|
|
// ODS-NEXT: let description = [{
|
|
// ODS-NEXT: It has two inputs.
|
|
// ODS-NEXT: It has one output.
|
|
// ODS-NEXT: }];
|
|
//
|
|
ods_def<Test6Op>:
|
|
def test6(A: f32(M, K), B: f32(K)) -> (C: f32(M))
|
|
"""
|
|
My magic op.
|
|
|
|
It has two inputs.
|
|
It has one output.
|
|
"""
|
|
{
|
|
C(m) = AddFOp<k>(C(m), MulFOp(A(m, k), B(k)));
|
|
}
|
|
|
|
// Test attribute builder
|
|
// ODS-LABEL: def Test7Op
|
|
// ODS: OpBuilder<
|
|
// ODS: (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
|
|
// ODS: "ValueRange":$outputs, "Attribute":$attr_a, "Attribute":$attr_b,
|
|
// ODS: CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)
|
|
// ODS: $_state.addAttribute("attr_a", attr_a);
|
|
// ODS: $_state.addAttribute("attr_b", attr_b);
|
|
//
|
|
ods_def<Test7Op>:
|
|
def test7(A: f32(M, K), B: f32(K)) -> (C: f32(M))
|
|
attr(attr_a: f32, attr_b: 4xi32)
|
|
{
|
|
C(m) = AddFOp<k>(C(m), MulFOp(A(m, k), B(k)));
|
|
}
|
|
|
|
// Test output arg order.
|
|
// IMPL-LABEL: void Test8Op::regionBuilder(ImplicitLocOpBuilder &b,
|
|
// IMPL: Block &block) {
|
|
// IMPL: Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
|
|
// IMPL: Value [[d:.*]] = b.create<MulFOp>([[a]], [[b]]);
|
|
// IMPL: Value [[e:.*]] = b.create<SubFOp>([[d]], [[c]]);
|
|
// IMPL: b.create<linalg::YieldOp>(ValueRange{ [[e]] });
|
|
ods_def<Test8Op>:
|
|
def test8(A: f32(M, K), B: f32(K)) -> (C: f32(M))
|
|
{
|
|
C(m) = SubFOp<k>(MulFOp(A(m, k), B(k)), C(m));
|
|
}
|
|
|
|
// Test shape-only operand.
|
|
// IMPL-LABEL: ArrayAttr Test9Op::indexing_maps() {
|
|
// IMPL: auto map0 = AffineMap::get(2, 2, {d0, d1}, context);
|
|
// IMPL: auto map1 = AffineMap::get(2, 2, {d1}, context);
|
|
// IMPL: auto map2 = AffineMap::get(2, 2, {d0}, context);
|
|
// IMPL-LABEL: void Test9Op::regionBuilder(ImplicitLocOpBuilder &b,
|
|
// IMPL: Block &block) {
|
|
// IMPL: Value [[a:.*]](args[0]), [[c:.*]](args[2]);
|
|
ods_def<Test9Op>:
|
|
def test9(A: f32(M, K), B: f32(K)) -> (C: f32(M))
|
|
{
|
|
C(m) = AddFOp<k>(C(m), A(m, k));
|
|
}
|