Skip to content

Commit 00104b5

Browse files
MaheshRavishankarMaheshRavishankar
authored
Allow dynamic dimensions during folding of tensor.expand_shape/collapse_shape into flow.dispatch.tensor.load/store. (#18873)
This also cleans up the implementation of these patterns to avoid using templated code that is hard to read/maintain. --------- Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>
1 parent a400cde commit 00104b5

File tree

3 files changed

+377
-72
lines changed

3 files changed

+377
-72
lines changed

‎compiler/src/iree/compiler/Codegen/Common/test/canonicalize_interface_load_store.mlir

+16-6
Original file line numberDiff line numberDiff line change
@@ -71,26 +71,36 @@ func.func @dont_fold_reshape_with_not_full_load() {
7171
// -----
7272

7373
#pipeline_layout = #hal.pipeline.layout<constants = 3, bindings = [
74+
#hal.pipeline.binding<storage_buffer>,
7475
#hal.pipeline.binding<storage_buffer>
7576
]>
76-
// CHECK-LABEL: func.func @dont_fold_dynamic_reshape()
77-
func.func @dont_fold_dynamic_reshape() {
77+
func.func @fold_dynamic_reshape() {
7878
%c0 = arith.constant 0 : index
7979
%c1 = arith.constant 1 : index
8080
%dim0 = hal.interface.constant.load layout(#pipeline_layout) ordinal(0) : index
8181
%dim1 = hal.interface.constant.load layout(#pipeline_layout) ordinal(1) : index
8282
%dim2 = hal.interface.constant.load layout(#pipeline_layout) ordinal(2) : index
8383
%1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) : !flow.dispatch.tensor<readonly:tensor<?x?x96xf32>>{%dim0, %dim1}
84-
%2 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) : !flow.dispatch.tensor<writeonly:tensor<?x12x8xf32>>{%dim2}
84+
%2 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) : !flow.dispatch.tensor<writeonly:tensor<?x12x8xf32>>{%dim2}
8585
%3 = flow.dispatch.tensor.load %1, offsets=[0, 0, 0], sizes =[%dim0, %dim1, 96], strides=[1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x96xf32>>{%dim0, %dim1} -> tensor<?x?x96xf32>
86-
// CHECK: tensor.collapse_shape
87-
// CHECK: tensor.expand_shape
8886
%4 = tensor.collapse_shape %3 [[0, 1], [2]] : tensor<?x?x96xf32> into tensor<?x96xf32>
8987
%dyn = tensor.dim %4, %c0 : tensor<?x96xf32>
9088
%5 = tensor.expand_shape %4 [[0], [1, 2]] output_shape [%dyn, 12, 8] : tensor<?x96xf32> into tensor<?x12x8xf32>
91-
flow.dispatch.tensor.store %5, %2, offsets = [%c0, %c0, %c0], sizes = [%c1, 12, 8], strides = [%c1, %c1, %c1] : tensor<?x12x8xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x12x8xf32>>{%dim2}
89+
flow.dispatch.tensor.store %5, %2, offsets = [0, 0, 0], sizes = [%dim2, 12, 8], strides = [1, 1, 1] : tensor<?x12x8xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x12x8xf32>>{%dim2}
9290
return
9391
}
92+
// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 * s1)>
93+
// CHECK: func.func @fold_dynamic_reshape()
94+
// CHECK-DAG: %[[CST0:.+]] = hal.interface.constant.load layout(#{{.+}}) ordinal(0)
95+
// CHECK-DAG: %[[CST1:.+]] = hal.interface.constant.load layout(#{{.+}}) ordinal(1)
96+
// CHECK-DAG: %[[CST2:.+]] = hal.interface.constant.load layout(#{{.+}}) ordinal(2)
97+
// CHECK: %[[COLLAPSED:.+]] = affine.apply #[[MAP]]()[%[[CST0]], %[[CST1]]]
98+
// CHECK: %[[IN_BINDING:.+]] = hal.interface.binding.subspan
99+
// CHECK-SAME: binding(0) : !flow.dispatch.tensor<readonly:tensor<?x96xf32>>{%[[COLLAPSED]]}
100+
// CHECK: %[[OUT_BINDING:.+]] = hal.interface.binding.subspan
101+
// CHECK-SAME: binding(1) : !flow.dispatch.tensor<writeonly:tensor<?x96xf32>>{%[[CST2]]}
102+
// CHECK: %[[IN:.+]] = flow.dispatch.tensor.load %[[IN_BINDING]]
103+
// CHECK: flow.dispatch.tensor.store %[[IN]], %[[OUT_BINDING]]
94104

95105
// -----
96106

‎compiler/src/iree/compiler/Codegen/Common/test/propagate_reshapes_by_expansion.mlir

+73-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-propagate-reshapes-by-expansion))" --split-input-file %s | FileCheck %s
1+
// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-propagate-reshapes-by-expansion))" --split-input-file %s --mlir-print-local-scope | FileCheck %s
22

33
func.func @reshape_and_lowering_config(%src: tensor<3x4xf16>, %dest: tensor<12xf16>, %dest2: tensor<12xf16>) -> tensor<12xf16> {
44
%collapse = tensor.collapse_shape %src [[0, 1]] : tensor<3x4xf16> into tensor<12xf16>
@@ -14,3 +14,75 @@ func.func @reshape_and_lowering_config(%src: tensor<3x4xf16>, %dest: tensor<12xf
1414
// CHECK: linalg.copy
1515
// CHECK-SAME: lowering_config = #iree_gpu.derived_thread_config
1616
// CHECK-SAME: ins(%[[COLLAPSE]]
17+
18+
// -----
19+
20+
#pipeline_layout = #hal.pipeline.layout<constants = 1, bindings = [
21+
#hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">], flags = Indirect>
22+
func.func @fold_collapse_into_loads_dynamic() -> tensor<?x32xf32> {
23+
%c0 = arith.constant 0 : index
24+
%0 = hal.interface.constant.load layout(#pipeline_layout) ordinal(0) : index
25+
%1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
26+
flags("ReadOnly|Indirect") : !flow.dispatch.tensor<readonly:tensor<2x?x32xf32>>{%0}
27+
%2 = flow.dispatch.tensor.load %1, offsets = [0, 0, 0], sizes = [2, %0, 32], strides = [1, 1, 1]
28+
: !flow.dispatch.tensor<readonly:tensor<2x?x32xf32>>{%0} -> tensor<2x?x32xf32>
29+
%3 = tensor.collapse_shape %2 [[0, 1], [2]] : tensor<2x?x32xf32> into tensor<?x32xf32>
30+
return %3 : tensor<?x32xf32>
31+
}
32+
// CHECK-LABEL: func @fold_collapse_into_loads_dynamic()
33+
// CHECK: %[[CONST:.+]] = hal.interface.constant.load
34+
// CHECK: %[[SHAPE:.+]] = affine.apply affine_map<()[s0] -> (s0 * 2)>()[%[[CONST]]]
35+
// CHECK: %[[SUBSPAN:.+]] = hal.interface.binding.subspan
36+
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x32xf32>>{%[[SHAPE]]}
37+
// CHECK: %[[LOAD:.+]] = flow.dispatch.tensor.load %[[SUBSPAN]]
38+
// CHECK-SAME: offsets = [0, 0], sizes = [%[[SHAPE]], 32], strides = [1, 1]
39+
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x32xf32>>{%[[SHAPE]]}
40+
41+
// -----
42+
43+
#pipeline_layout = #hal.pipeline.layout<constants = 2, bindings = [
44+
#hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">], flags = Indirect>
45+
func.func @fold_expand_into_loads_dynamic() -> tensor<2x?x16x32xf32> {
46+
%c0 = arith.constant 0 : index
47+
%0 = hal.interface.constant.load layout(#pipeline_layout) ordinal(0) : index
48+
%1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
49+
flags("ReadOnly|Indirect") : !flow.dispatch.tensor<readonly:tensor<2x?x32xf32>>{%0}
50+
%2 = flow.dispatch.tensor.load %1, offsets = [0, 0, 0], sizes = [2, %0, 32], strides = [1, 1, 1]
51+
: !flow.dispatch.tensor<readonly:tensor<2x?x32xf32>>{%0} -> tensor<2x?x32xf32>
52+
%3 = affine.apply affine_map<()[s0] -> (s0 floordiv 2)>()[%0]
53+
%4 = tensor.expand_shape %2 [[0], [1, 2], [3]] output_shape [2, %3, 16, 32] : tensor<2x?x32xf32> into tensor<2x?x16x32xf32>
54+
return %4 : tensor<2x?x16x32xf32>
55+
}
56+
// CHECK-LABEL: func @fold_expand_into_loads_dynamic()
57+
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
58+
// CHECK-DAG: %[[CONST:.+]] = hal.interface.constant.load
59+
// CHECK: %[[SHAPE:.+]] = arith.divui %[[CONST]], %[[C16]]
60+
// CHECK: %[[SUBSPAN:.+]] = hal.interface.binding.subspan
61+
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<2x?x16x32xf32>>{%[[SHAPE]]}
62+
// CHECK: %[[LOAD:.+]] = flow.dispatch.tensor.load %[[SUBSPAN]]
63+
// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [2, %[[SHAPE]], 16, 32], strides = [1, 1, 1, 1]
64+
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<2x?x16x32xf32>>{%[[SHAPE]]}
65+
66+
// -----
67+
68+
#pipeline_layout = #hal.pipeline.layout<constants = 1, bindings = [
69+
#hal.pipeline.binding<storage_buffer, Indirect>], flags = Indirect>
70+
func.func @fold_collapse_into_stores_dynamic(%arg0 : tensor<2x?x32xf32>) {
71+
%c0 = arith.constant 0 : index
72+
%0 = hal.interface.constant.load layout(#pipeline_layout) ordinal(0) : index
73+
%1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
74+
flags("ReadOnly|Indirect") : !flow.dispatch.tensor<writeonly:tensor<?x32xf32>>{%0}
75+
%2 = tensor.collapse_shape %arg0 [[0, 1], [2]] : tensor<2x?x32xf32> into tensor<?x32xf32>
76+
flow.dispatch.tensor.store %2, %1, offsets = [0, 0], sizes = [%0, 32], strides = [1, 1]
77+
: tensor<?x32xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x32xf32>>{%0}
78+
return
79+
}
80+
// CHECK-LABEL: func @fold_collapse_into_stores_dynamic(
81+
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
82+
// CHECK: %[[CONST:.+]] = hal.interface.constant.load
83+
// CHECK: %[[SHAPE:.+]] = arith.divui %[[CONST]], %[[C2]]
84+
// CHECK: %[[SUBSPAN:.+]] = hal.interface.binding.subspan
85+
// CHECK-SAME: !flow.dispatch.tensor<writeonly:tensor<2x?x32xf32>>{%[[SHAPE]]}
86+
// CHECK: flow.dispatch.tensor.store %{{.+}}, %[[SUBSPAN]]
87+
// CHECK-SAME: offsets = [0, 0, 0], sizes = [2, %[[SHAPE]], 32], strides = [1, 1, 1]
88+
// CHECK-SAME: !flow.dispatch.tensor<writeonly:tensor<2x?x32xf32>>{%[[SHAPE]]}

0 commit comments

Comments
 (0)