Skip to content

Commit 7a47113

Browse files
banach-spacebanach-space
authored
[mlir][nfc] Add tests for linalg.mmt4d (#81422)
linalg.mmt4d was added a while back (https://reviews.llvm.org/D105244), but there are virtually no tests in-tree. In the spirit of documenting through test, this PR adds a few basic examples.
1 parent 7a5c1a4 commit 7a47113

File tree

4 files changed

+132
-0
lines changed

4 files changed

+132
-0
lines changed

‎mlir/test/Dialect/Linalg/invalid.mlir

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -744,3 +744,29 @@ func.func @illegal_softmax_output_shape(%arg0: tensor<2x16x32xf32>) -> tensor<2x
744744
-> tensor<2x16xf32>
745745
return %1 : tensor<2x16xf32>
746746
}
747+
748+
// -----
749+
750+
func.func @mmt4d_dims_mismatch(%A: tensor<16x16x8x1xf32>,
751+
%B: tensor<16x16x8x1xf32>,
752+
%C_in: tensor<16x16x8x1xf32>) -> tensor<16x16x8x1xf32> {
753+
// expected-error @+1 {{inferred input/output operand #2 has shape's dimension #3 to be 8, but found 1}}
754+
%res = linalg.mmt4d
755+
ins(%A, %B: tensor<16x16x8x1xf32>, tensor<16x16x8x1xf32>)
756+
outs(%C_in: tensor<16x16x8x1xf32>)
757+
-> tensor<16x16x8x1xf32>
758+
return %res : tensor<16x16x8x1xf32>
759+
}
760+
761+
// -----
762+
763+
func.func @mmt4d_rank_mismatch(%A: tensor<16x16x8x1xf32>,
764+
%B: tensor<16x16x8x1xf32>,
765+
%C_in: tensor<8x8xf32>) -> tensor<8x8xf32> {
766+
// expected-error @+1 {{expected operand rank (2) to match the result rank of indexing_map #2 (4)}}
767+
%res = linalg.mmt4d
768+
ins(%A, %B: tensor<16x16x8x1xf32>, tensor<16x16x8x1xf32>)
769+
outs(%C_in: tensor<8x8xf32>)
770+
-> tensor<8x8xf32>
771+
return %res : tensor<8x8xf32>
772+
}

‎mlir/test/Dialect/Linalg/named-ops.mlir

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1219,6 +1219,17 @@ func.func @batchmatmul_transpose_b(%arg0: memref<2x3x5xf32>, %arg1: memref<2x7x5
12191219

12201220
// -----
12211221

1222+
// CHECK-LABEL: func @mmt4d
1223+
func.func @mmt4d(%A: tensor<10x32x8x1xf32>, %B: tensor<80x32x4x1xf32>, %C: tensor<10x80x8x4xf32>) -> tensor<10x80x8x4xf32> {
1224+
// CHECK: %{{.+}} = linalg.mmt4d
1225+
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<10x32x8x1xf32>, tensor<80x32x4x1xf32>)
1226+
// CHECK-SAME: outs(%{{.+}} : tensor<10x80x8x4xf32>) -> tensor<10x80x8x4xf32>
1227+
%0 = linalg.mmt4d ins(%A, %B : tensor<10x32x8x1xf32>, tensor<80x32x4x1xf32>) outs(%C: tensor<10x80x8x4xf32>) -> tensor<10x80x8x4xf32>
1228+
return %0: tensor<10x80x8x4xf32>
1229+
}
1230+
1231+
// -----
1232+
12221233
// CHECK-LABEL: func @batch_mmt4d
12231234
func.func @batch_mmt4d(%arg0: tensor<128x10x32x8x1xf32>, %arg1: tensor<128x80x32x4x1xf32>, %arg2: tensor<128x10x80x8x4xf32>) -> tensor<128x10x80x8x4xf32> {
12241235
// CHECK: %{{.+}} = linalg.batch_mmt4d

‎mlir/test/Dialect/Linalg/transform-op-mmt4d-to-fma.mlir

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
// RUN: mlir-opt %s -transform-interpreter | FileCheck %s
2+
3+
func.func @mmt4d_to_fma(%A: tensor<16x16x8x1xf32>, %B: tensor<16x16x8x1xf32>, %C_in: tensor<16x16x8x8xf32>) -> tensor<16x16x8x8xf32> {
4+
%res = linalg.mmt4d
5+
ins(%A, %B: tensor<16x16x8x1xf32>, tensor<16x16x8x1xf32>)
6+
outs(%C_in: tensor<16x16x8x8xf32>)
7+
-> tensor<16x16x8x8xf32>
8+
return %res : tensor<16x16x8x8xf32>
9+
}
10+
11+
12+
// CHECK-LABEL: @mmt4d_to_fma
13+
// CHECK-COUNT-8: vector.fma
14+
15+
module attributes {transform.with_named_sequence} {
16+
transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
17+
%func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.op<"func.func">
18+
19+
%mmt4d = transform.structured.match ops{["linalg.mmt4d"]} in %func
20+
21+
// Step 1: Tile
22+
: (!transform.op<"func.func">) -> !transform.any_op
23+
// Tile parallel dims
24+
%tiled_linalg_op_p, %loops:4 = transform.structured.tile_using_for %mmt4d[1, 1, 0, 8, 8, 0]
25+
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
26+
// Tile reduction dims
27+
%tiled_linalg_op_r, %loops2:2 = transform.structured.tile_using_for %tiled_linalg_op_p[0, 0, 1, 0, 0, 1]
28+
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
29+
30+
// Step 2: Vectorize
31+
transform.structured.vectorize %tiled_linalg_op_r : !transform.any_op
32+
33+
// Step 3: Simplify
34+
// vector.multi_reduction --> vector.contract
35+
// Generates a 6-dim vector.contract with the dim matching the original MMT4D Op
36+
// and with the following split into parallel and reduction dims:
37+
// * parallel, parallel, reduction, parallel, parallel, reduction
38+
transform.apply_patterns to %func {
39+
transform.apply_patterns.vector.reduction_to_contract
40+
// Reduce the rank of xfer ops. This transforms vector.contract to be
41+
// more matmul-like and to enable the lowering to outer product Ops.
42+
transform.apply_patterns.vector.transfer_permutation_patterns
43+
} : !transform.op<"func.func">
44+
45+
// Hoisting and LICM - not strictly required
46+
%func_h = transform.structured.hoist_redundant_vector_transfers %func
47+
: (!transform.op<"func.func">) -> !transform.op<"func.func">
48+
%all_loops = transform.structured.match interface{LoopLikeInterface} in %func_h
49+
: (!transform.op<"func.func">) -> !transform.any_op
50+
transform.apply_licm to %all_loops : !transform.any_op
51+
transform.loop.hoist_loop_invariant_subsets %all_loops : !transform.any_op
52+
53+
// Simplify the 6-dim vector.contract into a 3-dim matmul-like
54+
// vector.contract with the following split into parallel and reduction
55+
// dims:
56+
// * parallel, parallel, reduction
57+
transform.apply_patterns to %func_h {
58+
transform.apply_patterns.vector.reduction_to_contract
59+
transform.apply_patterns.vector.cast_away_vector_leading_one_dim
60+
transform.apply_patterns.canonicalization
61+
} : !transform.op<"func.func">
62+
63+
// Step 4: Lower vector.contract to vector.fma via vector.outerproduct
64+
transform.apply_patterns to %func_h {
65+
transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
66+
transform.apply_patterns.vector.lower_outerproduct
67+
} : !transform.op<"func.func">
68+
transform.yield
69+
}
70+
}

‎mlir/test/Dialect/Linalg/vectorization.mlir

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -639,6 +639,31 @@ module attributes {transform.with_named_sequence} {
639639

640640
// -----
641641

642+
func.func @mmt4d(%A: memref<16x16x8x1xf32>, %B: memref<16x16x8x1xf32>, %C_in: memref<16x16x8x8xf32>) {
643+
linalg.mmt4d ins(%A, %B: memref<16x16x8x1xf32>, memref<16x16x8x1xf32>)
644+
outs(%C_in: memref<16x16x8x8xf32>)
645+
return
646+
}
647+
648+
// CHECK-LABEL: func.func @mmt4d(
649+
// CHECK-SAME: %[[A:.*]]: memref<16x16x8x1xf32>, %[[B:.*]]: memref<16x16x8x1xf32>, %[[C:.*]]: memref<16x16x8x8xf32>) {
650+
// CHECK: %[[VEC_A:.*]] = vector.transfer_read %[[A]]{{.*}} : memref<16x16x8x1xf32>, vector<16x16x16x8x8x1xf32>
651+
// CHECK: %[[VEC_B:.*]] = vector.transfer_read %[[B]]{{.*}} : memref<16x16x8x1xf32>, vector<16x16x16x8x8x1xf32>
652+
// CHECK: %[[VEC_C:.*]] = vector.transfer_read %[[C]]{{.*}} : memref<16x16x8x8xf32>, vector<16x16x8x8xf32>
653+
// CHECK: %[[MUL:.*]] = arith.mulf %[[VEC_A]], %[[VEC_B]] : vector<16x16x16x8x8x1xf32>
654+
// CHECK: %[[RED:.*]] = vector.multi_reduction <add>, %[[MUL]], %[[VEC_C]] [2, 5] : vector<16x16x16x8x8x1xf32> to vector<16x16x8x8xf32>
655+
// CHECK: vector.transfer_write %[[RED]], %[[C]]{{.*}} : vector<16x16x8x8xf32>, memref<16x16x8x8xf32>
656+
657+
module attributes {transform.with_named_sequence} {
658+
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
659+
%mmt4d = transform.structured.match ops{["linalg.mmt4d"]} in %arg1 : (!transform.any_op) -> !transform.any_op
660+
transform.structured.vectorize %mmt4d : !transform.any_op
661+
transform.yield
662+
}
663+
}
664+
665+
// -----
666+
642667
func.func @matmul_scalable(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {
643668
linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)
644669
outs(%C: memref<?x?xf32>)

0 commit comments

Comments
 (0)