Skip to content

Commit b32b16a

Browse files
oowekyalaoowekyala
committed
Update tests
1 parent bc36549 commit b32b16a

File tree

2 files changed

+80
-75
lines changed

2 files changed

+80
-75
lines changed

‎test/Dialect/Cinm/cinm-tiling.mlir‎

Lines changed: 59 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -1,70 +1,72 @@
1-
// RUN: cinm-opt %s --cinm-tiling -split-input-file | FileCheck %s
1+
// // RUN: cinm-opt %s --cinm-tiling -split-input-file | FileCheck %s
2+
// this file is old i think, i meant to replace it with cinm-tiling2
3+
// todo add test for tiling of elementwise in the other file
24

35

4-
// CHECK-LABEL: @gemmSquare
5-
// CHECK-SAME: (%[[A:.*]]: tensor<1024x1024xi32>, %[[B:.*]]: tensor<1024x1024xi32>) -> tensor<1024x1024xi32> {
6-
// CHECK: %[[res0:.*]] = affine.for %[[i:.*]] = 0 to 1024 iter_args({{.*}})
7-
// CHECK-NEXT: %[[res1:.*]] = affine.for %[[j:.*]] = 0 to 1024 step 1024 iter_args(%[[acc0:.*]] = {{.*}})
8-
// CHECK: %[[res2:.*]] = affine.for %[[k:.*]] = 0 to 1024 step 256 iter_args(%[[acc1:.*]] = {{.*}})
9-
// CHECK-NEXT: %[[blockA:.*]] = tensor.extract_slice %[[A]][%[[i]], %[[k]]] [1, 256] [1, 1]
10-
// CHECK-NEXT: %[[blockB:.*]] = tensor.extract_slice %[[B]][%[[k]], %[[j]]] [256, 1024] [1, 1]
11-
// CHECK-NEXT: %[[res3:.*]] = cinm.op.gemm %[[blockA]], %[[blockB]] plus %[[acc1]] {cinm.notile}
12-
// CHECK-NEXT: affine.yield %[[res3]] : tensor<1x1024xi32>
13-
// CHECK: %[[ins:.*]] = tensor.insert_slice %[[res2]] into %[[acc0]][%[[i]], %[[j]]]
14-
// CHECK-NEXT: affine.yield %[[ins]] : tensor<1024x1024xi32>
15-
// CHECK: affine.yield %[[res1]] : tensor<1024x1024xi32>
16-
// CHECK: cinm.yield %[[res0]] : tensor<1024x1024xi32>
6+
// // CHECK-LABEL: @gemmSquare
7+
// // CHECK-SAME: (%[[A:.*]]: tensor<1024x1024xi32>, %[[B:.*]]: tensor<1024x1024xi32>) -> tensor<1024x1024xi32> {
8+
// // CHECK: %[[res0:.*]] = affine.for %[[i:.*]] = 0 to 1024 iter_args({{.*}})
9+
// // CHECK-NEXT: %[[res1:.*]] = affine.for %[[j:.*]] = 0 to 1024 step 1024 iter_args(%[[acc0:.*]] = {{.*}})
10+
// // CHECK: %[[res2:.*]] = affine.for %[[k:.*]] = 0 to 1024 step 256 iter_args(%[[acc1:.*]] = {{.*}})
11+
// // CHECK-NEXT: %[[blockA:.*]] = tensor.extract_slice %[[A]][%[[i]], %[[k]]] [1, 256] [1, 1]
12+
// // CHECK-NEXT: %[[blockB:.*]] = tensor.extract_slice %[[B]][%[[k]], %[[j]]] [256, 1024] [1, 1]
13+
// // CHECK-NEXT: %[[res3:.*]] = cinm.op.gemm %[[blockA]], %[[blockB]] plus %[[acc1]] {cinm.notile}
14+
// // CHECK-NEXT: affine.yield %[[res3]] : tensor<1x1024xi32>
15+
// // CHECK: %[[ins:.*]] = tensor.insert_slice %[[res2]] into %[[acc0]][%[[i]], %[[j]]]
16+
// // CHECK-NEXT: affine.yield %[[ins]] : tensor<1024x1024xi32>
17+
// // CHECK: affine.yield %[[res1]] : tensor<1024x1024xi32>
18+
// // CHECK: cinm.yield %[[res0]] : tensor<1024x1024xi32>
1719

18-
func.func @gemmSquare(%a: tensor<1024x1024xi32>, %b: tensor<1024x1024xi32>) -> tensor<1024x1024xi32> {
19-
%res = cinm.compute attributes { workgroupShape = array<i64: 4>, bufferSizesInBytes = array<i64: 1024> } -> tensor<1024x1024xi32> {
20-
%d = cinm.op.gemm %a, %b : tensor<1024x1024xi32>, tensor<1024x1024xi32> -> tensor<1024x1024xi32>
21-
cinm.yield %d: tensor<1024x1024xi32>
22-
}
23-
return %res: tensor<1024x1024xi32>
24-
}
20+
// func.func @gemmSquare(%a: tensor<1024x1024xi32>, %b: tensor<1024x1024xi32>) -> tensor<1024x1024xi32> {
21+
// %res = cinm.compute (%a0 = %a: tensor<1024x1024xi32>, %b0 = %b: tensor<1024x1024xi32>) -> tensor<1024x1024xi32> attributes { workgroupShape = array<i64: 4>, bufferSizesInBytes = array<i64: 1024> } {
22+
// %d = cinm.op.gemm %a0, %b0 : tensor<1024x1024xi32>, tensor<1024x1024xi32> -> tensor<1024x1024xi32>
23+
// cinm.yield %d: tensor<1024x1024xi32>
24+
// }
25+
// return %res: tensor<1024x1024xi32>
26+
// }
2527

2628

27-
// -----
29+
// // -----
2830

29-
// CHECK-LABEL: @gemv
31+
// // CHECK-LABEL: @gemv
3032

31-
func.func @gemv(%a: tensor<1024x1024xi32>, %b: tensor<1024xi32>) -> tensor<1024xi32>{
32-
%res = cinm.compute attributes { workgroupShape = array<i64: 4>, bufferSizesInBytes = array<i64: 1024> }-> tensor<1024xi32> {
33-
%d = cinm.op.gemv %a, %b : tensor<1024x1024xi32>, tensor<1024xi32> -> tensor<1024xi32>
34-
cinm.yield %d: tensor<1024xi32>
35-
}
36-
return %res: tensor<1024xi32>
37-
}
33+
// func.func @gemv(%a: tensor<1024x1024xi32>, %b: tensor<1024xi32>) -> tensor<1024xi32>{
34+
// %res = cinm.compute (%a0 = %a: tensor<1024x1024xi32>, %b0 = %b: tensor<1024xi32>) -> tensor<1024xi32> attributes { workgroupShape = array<i64: 4>, bufferSizesInBytes = array<i64: 1024> } {
35+
// %d = cinm.op.gemv %a0, %b0 : tensor<1024x1024xi32>, tensor<1024xi32> -> tensor<1024xi32>
36+
// cinm.yield %d: tensor<1024xi32>
37+
// }
38+
// return %res: tensor<1024xi32>
39+
// }
3840

39-
// -----
41+
// // -----
4042

41-
// CHECK-LABEL: @max
42-
// CHECK-SAME: (%[[input:.*]]: tensor<1024xi32>) -> i32
43-
// CHECK-NEXT: %[[res:.*]] = cinm.compute attributes {{{.*}}} -> i32 {
44-
// CHECK: %[[gen:.*]] = tensor.generate {
45-
// CHECK-NEXT: ^{{.*}}(%[[idx:.*]]: {{.*}}):
46-
// CHECK-NEXT: %[[idxOffset:.*]] = arith.muli %[[idx]]
47-
// CHECK-NEXT: %[[extracted:.*]] = tensor.extract_slice %[[input]][%[[idx]]] [256] [1]
48-
// CHECK-NEXT: %[[redInner:.*]] = linalg.reduce ins(%[[extracted]] : {{.*}}) outs({{.*}}) dimensions = [0]
49-
// CHECK-NEXT: (%[[in0:.*]]: {{.*}}, %[[acc0:.*]]: {{.*}})
50-
// CHECK-NEXT: %[[res0:.*]] = arith.maxsi %[[in0]], %[[acc0]]
51-
// CHECK-NEXT: linalg.yield %[[res0]]
43+
// // CHECK-LABEL: @max
44+
// // CHECK-SAME: (%[[input:.*]]: tensor<1024xi32>) -> i32
45+
// // CHECK-NEXT: %[[res:.*]] = cinm.compute attributes {{{.*}}} -> i32 {
46+
// // CHECK: %[[gen:.*]] = tensor.generate {
47+
// // CHECK-NEXT: ^{{.*}}(%[[idx:.*]]: {{.*}}):
48+
// // CHECK-NEXT: %[[idxOffset:.*]] = arith.muli %[[idx]]
49+
// // CHECK-NEXT: %[[extracted:.*]] = tensor.extract_slice %[[input]][%[[idx]]] [256] [1]
50+
// // CHECK-NEXT: %[[redInner:.*]] = linalg.reduce ins(%[[extracted]] : {{.*}}) outs({{.*}}) dimensions = [0]
51+
// // CHECK-NEXT: (%[[in0:.*]]: {{.*}}, %[[acc0:.*]]: {{.*}})
52+
// // CHECK-NEXT: %[[res0:.*]] = arith.maxsi %[[in0]], %[[acc0]]
53+
// // CHECK-NEXT: linalg.yield %[[res0]]
5254

53-
// CHECK: %[[extracted0:.*]] = tensor.extract %[[redInner]][] : tensor<i32>
54-
// CHECK-NEXT: tensor.yield %[[extracted0]]
55+
// // CHECK: %[[extracted0:.*]] = tensor.extract %[[redInner]][] : tensor<i32>
56+
// // CHECK-NEXT: tensor.yield %[[extracted0]]
5557

56-
// CHECK: %[[redOuter:.*]] = linalg.reduce ins(%[[gen]] : tensor<4xi32>) outs({{.*}}) dimensions = [0]
57-
// CHECK-NEXT: (%[[in1:.*]]: {{.*}}, %[[acc1:.*]]: {{.*}})
58-
// CHECK-NEXT: %[[res1:.*]] = arith.maxsi %[[in1]], %[[acc1]]
59-
// CHECK-NEXT: linalg.yield %[[res1]]
58+
// // CHECK: %[[redOuter:.*]] = linalg.reduce ins(%[[gen]] : tensor<4xi32>) outs({{.*}}) dimensions = [0]
59+
// // CHECK-NEXT: (%[[in1:.*]]: {{.*}}, %[[acc1:.*]]: {{.*}})
60+
// // CHECK-NEXT: %[[res1:.*]] = arith.maxsi %[[in1]], %[[acc1]]
61+
// // CHECK-NEXT: linalg.yield %[[res1]]
6062

61-
// CHECK: %[[extracted1:.*]] = tensor.extract %[[redOuter]][]
62-
// CHECK-NEXT: cinm.yield %[[extracted1]]
63+
// // CHECK: %[[extracted1:.*]] = tensor.extract %[[redOuter]][]
64+
// // CHECK-NEXT: cinm.yield %[[extracted1]]
6365

64-
func.func @max(%a: tensor<1024xi32>) -> i32 {
65-
%res = cinm.compute attributes { workgroupShape = array<i64: 4>, bufferSizesInBytes = array<i64: 1024> } -> i32 {
66-
%d = cinm.op.reduce max (%a): tensor<1024xi32> -> i32
67-
cinm.yield %d : i32
68-
}
69-
return %res: i32
70-
}
66+
// func.func @max(%a: tensor<1024xi32>) -> i32 {
67+
// %res = cinm.compute (%a0 = %a : tensor<1024xi32>) -> i32 attributes { workgroupShape = array<i64: 4>, bufferSizesInBytes = array<i64: 1024> } {
68+
// %d = cinm.op.reduce max (%a0): tensor<1024xi32> -> i32
69+
// cinm.yield %d : i32
70+
// }
71+
// return %res: i32
72+
// }

‎test/Dialect/Cinm/cinm-tiling2.mlir‎

Lines changed: 21 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,21 @@
33
// CHECK-LABEL: @gemm_memref
44

55
func.func @gemm_memref(%arg0: memref<8x1024xi32>, %arg1: memref<1024x128xi32>) -> memref<8x128xi32> {
6+
// CHECK: cinm.compute (%[[a0:.*]] = %{{.*}}, %[[b0:.*]] = %{{.*}}) ->
67
// CHECK: %[[out:.*]] = memref.alloc()
78
// CHECK: linalg.fill ins({{.*}}) outs(%[[out]] :
89
// CHECK: affine.for %[[i:.*]] = 0 to 8 step 8
910
// CHECK: affine.for %[[j:.*]] = 0 to 128 step 128
1011
// CHECK: %[[sliceOut:.*]] = memref.subview %[[out]][%[[i]], %[[j]]] [8, 128] [1, 1] :
1112
// CHECK: affine.for %[[k:.*]] = 0 to 1024 step 32
12-
// CHECK: %[[sliceA:.*]] = memref.subview %arg0[%[[i]], %[[k]]] [8, 32] [1, 1] :
13-
// CHECK: %[[sliceB:.*]] = memref.subview %arg1[%[[k]], %[[j]]] [32, 128] [1, 1] :
13+
// CHECK: %[[sliceA:.*]] = memref.subview %[[a0]][%[[i]], %[[k]]] [8, 32] [1, 1] :
14+
// CHECK: %[[sliceB:.*]] = memref.subview %[[b0]][%[[k]], %[[j]]] [32, 128] [1, 1] :
1415
// CHECK: cinm.op.gemm %[[sliceA]], %[[sliceB]] into %[[sliceOut]] {cinm.notile} :
15-
%0 = cinm.compute attributes {workgroupShape = array<i64: 8, 128, 1>, bufferSizesInBytes=array<i64: 0,0,512>} -> memref<8x128xi32> {
16+
%0 = cinm.compute (%a0 = %arg0 : memref<8x1024xi32>, %a1 = %arg1: memref<1024x128xi32>) -> memref<8x128xi32> attributes {workgroupShape = array<i64: 8, 128, 1>, bufferSizesInBytes=array<i64: 0,0,512>} {
1617
%alloc = memref.alloc() : memref<8x128xi32>
1718
%c0_i32 = arith.constant 0 : i32
1819
linalg.fill ins(%c0_i32 : i32) outs(%alloc : memref<8x128xi32>)
19-
cinm.op.gemm %arg0, %arg1 into %alloc : memref<8x1024xi32>, memref<1024x128xi32> into memref<8x128xi32>
20+
cinm.op.gemm %a0, %a1 into %alloc : memref<8x1024xi32>, memref<1024x128xi32> into memref<8x128xi32>
2021
cinm.yield %alloc : memref<8x128xi32>
2122
}
2223
return %0 : memref<8x128xi32>
@@ -27,62 +28,64 @@ func.func @gemm_memref(%arg0: memref<8x1024xi32>, %arg1: memref<1024x128xi32>) -
2728
// CHECK-SAME: ({{.*}}, %[[bias:.*]]: memref<8x128xi32>)
2829

2930
func.func @gemm_memref_bias(%arg0: memref<8x1024xi32>, %arg1: memref<1024x128xi32>, %bias: memref<8x128xi32>) -> memref<8x128xi32> {
31+
// CHECK: cinm.compute (%[[a0:.*]] = %{{.*}}, %[[b0:.*]] = %{{.*}}, %[[c0:.*]] = %{{.*}}) ->
3032
// CHECK: %[[out:.*]] = memref.alloc()
3133
// CHECK: affine.for %[[i:.*]] = 0 to 8 step 8
3234
// CHECK: affine.for %[[j:.*]] = 0 to 128 step 128
33-
// CHECK: %[[sliceBias:.*]] = memref.subview %arg2[%[[i]], %[[j]]] [8, 128] [1, 1] :
35+
// CHECK: %[[sliceBias:.*]] = memref.subview %[[c0]][%[[i]], %[[j]]] [8, 128] [1, 1] :
3436
// CHECK: %[[sliceOut:.*]] = memref.subview %[[out]][%[[i]], %[[j]]] [8, 128] [1, 1] :
3537
// CHECK: linalg.add ins(%[[sliceBias]], %[[sliceOut]] : {{.*}}) outs(%[[sliceOut]] :
3638
// CHECK: affine.for %[[k:.*]] = 0 to 1024 step 32
37-
// CHECK: %[[sliceA:.*]] = memref.subview %arg0[%[[i]], %[[k]]] [8, 32] [1, 1] :
38-
// CHECK: %[[sliceB:.*]] = memref.subview %arg1[%[[k]], %[[j]]] [32, 128] [1, 1] :
39+
// CHECK: %[[sliceA:.*]] = memref.subview %[[a0]][%[[i]], %[[k]]] [8, 32] [1, 1] :
40+
// CHECK: %[[sliceB:.*]] = memref.subview %[[b0]][%[[k]], %[[j]]] [32, 128] [1, 1] :
3941
// CHECK: cinm.op.gemm %[[sliceA]], %[[sliceB]] into %[[sliceOut]] {cinm.notile} :
40-
%0 = cinm.compute attributes {workgroupShape = array<i64: 8, 128, 1>, bufferSizesInBytes=array<i64: 0,0,512>} -> memref<8x128xi32> {
42+
%0 = cinm.compute(%a0 = %arg0: memref<8x1024xi32>, %a1 = %arg1: memref<1024x128xi32>, %b0 = %bias: memref<8x128xi32>) -> memref<8x128xi32> attributes {workgroupShape = array<i64: 8, 128, 1>, bufferSizesInBytes=array<i64: 0,0,512>} {
4143
%alloc = memref.alloc() : memref<8x128xi32>
4244
%c0_i32 = arith.constant 0 : i32
4345
linalg.fill ins(%c0_i32 : i32) outs(%alloc : memref<8x128xi32>)
44-
cinm.op.gemm %arg0, %arg1 plus %bias into %alloc : memref<8x1024xi32>, memref<1024x128xi32> plus memref<8x128xi32> into memref<8x128xi32>
46+
cinm.op.gemm %a0, %a1 plus %b0 into %alloc : memref<8x1024xi32>, memref<1024x128xi32> plus memref<8x128xi32> into memref<8x128xi32>
4547
cinm.yield %alloc : memref<8x128xi32>
4648
}
4749
return %0 : memref<8x128xi32>
4850
}
4951

5052
// -----
5153
// CHECK-LABEL: @gemm_tensor
54+
// CHECK: cinm.compute (%[[a0:.*]] = %{{.*}}, %[[b0:.*]] = %{{.*}}) ->
5255
// CHECK: affine.for %[[i:.*]] = 0 to 8 step 8 iter_args(%
5356
// CHECK: affine.for %[[j:.*]] = 0 to 128 step 128 iter_args(%[[outer:.*]] =
5457
// CHECK: %[[innerinit:.*]] = arith.constant dense<0> :
5558
// CHECK: %[[x:.*]] = affine.for %[[k:.*]] = 0 to 1024 step 32 iter_args(%[[inner:.*]] = %[[innerinit]])
5659

57-
// CHECK: %[[sliceA:.*]] = tensor.extract_slice %arg0[%[[i]], %[[k]]] [8, 32] [1, 1] :
58-
// CHECK: %[[sliceB:.*]] = tensor.extract_slice %arg1[%[[k]], %[[j]]] [32, 128] [1, 1] :
60+
// CHECK: %[[sliceA:.*]] = tensor.extract_slice %[[a0]][%[[i]], %[[k]]] [8, 32] [1, 1] :
61+
// CHECK: %[[sliceB:.*]] = tensor.extract_slice %[[b0]][%[[k]], %[[j]]] [32, 128] [1, 1] :
5962
// CHECK: %[[r:.*]] = cinm.op.gemm %[[sliceA]], %[[sliceB]] plus %[[inner]] {cinm.notile} :
6063
// CHECK: affine.yield %[[r]]
6164
// CHECK: tensor.insert_slice %[[x]] into %[[outer]][%[[i]], %[[j]]] [8, 128] [1, 1] :
6265
func.func @gemm_tensor(%A: tensor<8x1024xi32>, %B: tensor<1024x128xi32>) -> tensor<8x128xi32> {
63-
%r0 = cinm.compute attributes { workgroupShape=array<i64: 8, 128, 1>, bufferSizesInBytes=array<i64: 0,0,512> } -> tensor<8x128xi32> {
64-
%r = cinm.op.gemm %A, %B: tensor<8x1024xi32>, tensor<1024x128xi32> -> tensor<8x128xi32>
66+
%r0 = cinm.compute (%a = %A: tensor<8x1024xi32>, %b = %B: tensor<1024x128xi32>) -> tensor<8x128xi32> attributes { workgroupShape=array<i64: 8, 128, 1>, bufferSizesInBytes=array<i64: 0,0,512> } {
67+
%r = cinm.op.gemm %a, %b: tensor<8x1024xi32>, tensor<1024x128xi32> -> tensor<8x128xi32>
6568
cinm.yield %r : tensor<8x128xi32>
6669
}
6770
func.return %r0 : tensor<8x128xi32>
6871
}
6972

7073
// -----
7174
// CHECK-LABEL: @gemm_tensor_bias
72-
// CHECK-SAME: ({{.*}}, %[[bias:.*]]: tensor<8x128xi32>)
75+
// CHECK: cinm.compute (%[[a0:.*]] = %{{.*}}, %[[b0:.*]] = %{{.*}}, %[[bias:.*]] = %{{.*}}) ->
7376
// CHECK: affine.for %[[i:.*]] = 0 to 8 step 8 iter_args(%
7477
// CHECK: affine.for %[[j:.*]] = 0 to 128 step 128 iter_args(%[[outer:.*]] =
7578
// CHECK: %[[innerinit:.*]] = tensor.extract_slice %[[bias]][%[[i]], %[[j]]] [8, 128] [1, 1] :
7679
// CHECK: %[[x:.*]] = affine.for %[[k:.*]] = 0 to 1024 step 32 iter_args(%[[inner:.*]] = %[[innerinit]])
7780

78-
// CHECK: %[[sliceA:.*]] = tensor.extract_slice %arg0[%[[i]], %[[k]]] [8, 32] [1, 1] :
79-
// CHECK: %[[sliceB:.*]] = tensor.extract_slice %arg1[%[[k]], %[[j]]] [32, 128] [1, 1] :
81+
// CHECK: %[[sliceA:.*]] = tensor.extract_slice %[[a0]][%[[i]], %[[k]]] [8, 32] [1, 1] :
82+
// CHECK: %[[sliceB:.*]] = tensor.extract_slice %[[b0]][%[[k]], %[[j]]] [32, 128] [1, 1] :
8083
// CHECK: %[[r:.*]] = cinm.op.gemm %[[sliceA]], %[[sliceB]] plus %[[inner]] {cinm.notile} :
8184
// CHECK: affine.yield %[[r]]
8285
// CHECK: tensor.insert_slice %[[x]] into %[[outer]][%[[i]], %[[j]]] [8, 128] [1, 1] :
8386
func.func @gemm_tensor_bias(%A: tensor<8x1024xi32>, %B: tensor<1024x128xi32>, %bias: tensor<8x128xi32>) -> tensor<8x128xi32> {
84-
%r0 = cinm.compute attributes { workgroupShape=array<i64: 8, 128, 1>, bufferSizesInBytes=array<i64: 0,0,512> } -> tensor<8x128xi32> {
85-
%r = cinm.op.gemm %A, %B plus %bias: tensor<8x1024xi32>, tensor<1024x128xi32> plus tensor<8x128xi32> -> tensor<8x128xi32>
87+
%r0 = cinm.compute(%a = %A: tensor<8x1024xi32>, %b = %B: tensor<1024x128xi32>, %c = %bias: tensor<8x128xi32>) -> tensor<8x128xi32> attributes { workgroupShape=array<i64: 8, 128, 1>, bufferSizesInBytes=array<i64: 0,0,512> } {
88+
%r = cinm.op.gemm %a, %b plus %c: tensor<8x1024xi32>, tensor<1024x128xi32> plus tensor<8x128xi32> -> tensor<8x128xi32>
8689
cinm.yield %r : tensor<8x128xi32>
8790
}
8891
func.return %r0 : tensor<8x128xi32>

0 commit comments

Comments
 (0)