legalization.mlir
3.81 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
// RUN: mlir-opt -legalize-std-for-spirv -verify-diagnostics %s -o - | FileCheck %s
// CHECK-LABEL: @fold_static_stride_subview_with_load
// CHECK-SAME: [[ARG0:%.*]]: memref<12x32xf32>, [[ARG1:%.*]]: index, [[ARG2:%.*]]: index, [[ARG3:%.*]]: index, [[ARG4:%.*]]: index
func @fold_static_stride_subview_with_load(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index) {
// CHECK-NOT: subview
// CHECK: [[C2:%.*]] = constant 2 : index
// CHECK: [[C3:%.*]] = constant 3 : index
// CHECK: [[STRIDE1:%.*]] = muli [[ARG3]], [[C2]] : index
// CHECK: [[INDEX1:%.*]] = addi [[ARG1]], [[STRIDE1]] : index
// CHECK: [[STRIDE2:%.*]] = muli [[ARG4]], [[C3]] : index
// CHECK: [[INDEX2:%.*]] = addi [[ARG2]], [[STRIDE2]] : index
// CHECK: load [[ARG0]]{{\[}}[[INDEX1]], [[INDEX2]]{{\]}}
%0 = subview %arg0[%arg1, %arg2][][] : memref<12x32xf32> to memref<4x4xf32, offset:?, strides: [64, 3]>
%1 = load %0[%arg3, %arg4] : memref<4x4xf32, offset:?, strides: [64, 3]>
return
}
// CHECK-LABEL: @fold_dynamic_stride_subview_with_load
// CHECK-SAME: [[ARG0:%.*]]: memref<12x32xf32>, [[ARG1:%.*]]: index, [[ARG2:%.*]]: index, [[ARG3:%.*]]: index, [[ARG4:%.*]]: index, [[ARG5:%.*]]: index, [[ARG6:%.*]]: index
func @fold_dynamic_stride_subview_with_load(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index) {
// CHECK-NOT: subview
// CHECK: [[STRIDE1:%.*]] = muli [[ARG3]], [[ARG5]] : index
// CHECK: [[INDEX1:%.*]] = addi [[ARG1]], [[STRIDE1]] : index
// CHECK: [[STRIDE2:%.*]] = muli [[ARG4]], [[ARG6]] : index
// CHECK: [[INDEX2:%.*]] = addi [[ARG2]], [[STRIDE2]] : index
// CHECK: load [[ARG0]]{{\[}}[[INDEX1]], [[INDEX2]]{{\]}}
%0 = subview %arg0[%arg1, %arg2][][%arg5, %arg6] : memref<12x32xf32> to memref<4x4xf32, offset:?, strides: [?, ?]>
%1 = load %0[%arg3, %arg4] : memref<4x4xf32, offset:?, strides: [?, ?]>
return
}
// CHECK-LABEL: @fold_static_stride_subview_with_store
// CHECK-SAME: [[ARG0:%.*]]: memref<12x32xf32>, [[ARG1:%.*]]: index, [[ARG2:%.*]]: index, [[ARG3:%.*]]: index, [[ARG4:%.*]]: index, [[ARG5:%.*]]: f32
func @fold_static_stride_subview_with_store(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : f32) {
// CHECK-NOT: subview
// CHECK: [[C2:%.*]] = constant 2 : index
// CHECK: [[C3:%.*]] = constant 3 : index
// CHECK: [[STRIDE1:%.*]] = muli [[ARG3]], [[C2]] : index
// CHECK: [[INDEX1:%.*]] = addi [[ARG1]], [[STRIDE1]] : index
// CHECK: [[STRIDE2:%.*]] = muli [[ARG4]], [[C3]] : index
// CHECK: [[INDEX2:%.*]] = addi [[ARG2]], [[STRIDE2]] : index
// CHECK: store [[ARG5]], [[ARG0]]{{\[}}[[INDEX1]], [[INDEX2]]{{\]}}
%0 = subview %arg0[%arg1, %arg2][][] : memref<12x32xf32> to memref<4x4xf32, offset:?, strides: [64, 3]>
store %arg5, %0[%arg3, %arg4] : memref<4x4xf32, offset:?, strides: [64, 3]>
return
}
// CHECK-LABEL: @fold_dynamic_stride_subview_with_store
// CHECK-SAME: [[ARG0:%.*]]: memref<12x32xf32>, [[ARG1:%.*]]: index, [[ARG2:%.*]]: index, [[ARG3:%.*]]: index, [[ARG4:%.*]]: index, [[ARG5:%.*]]: index, [[ARG6:%.*]]: index, [[ARG7:%.*]]: f32
func @fold_dynamic_stride_subview_with_store(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index, %arg7 : f32) {
// CHECK-NOT: subview
// CHECK: [[STRIDE1:%.*]] = muli [[ARG3]], [[ARG5]] : index
// CHECK: [[INDEX1:%.*]] = addi [[ARG1]], [[STRIDE1]] : index
// CHECK: [[STRIDE2:%.*]] = muli [[ARG4]], [[ARG6]] : index
// CHECK: [[INDEX2:%.*]] = addi [[ARG2]], [[STRIDE2]] : index
// CHECK: store [[ARG7]], [[ARG0]]{{\[}}[[INDEX1]], [[INDEX2]]{{\]}}
%0 = subview %arg0[%arg1, %arg2][][%arg5, %arg6] : memref<12x32xf32> to memref<4x4xf32, offset:?, strides: [?, ?]>
store %arg7, %0[%arg3, %arg4] : memref<4x4xf32, offset:?, strides: [?, ?]>
return
}