Skip to content

Commit

Permalink
Allow dynamic dimensions during folding of `tensor.expand_shape/colla…
Browse files Browse the repository at this point in the history 
…pse_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 <[email protected]>
  • Loading branch information
@MaheshRavishankar
MaheshRavishankar authored Oct 23, 2024
1 parent a400cde commit 00104b5
Show file tree
Hide file tree
Showing 3 changed files with 377 additions and 72 deletions.
22 changes: 16 additions & 6 deletions compiler/src/iree/compiler/Codegen/Common/test/canonicalize_interface_load_store.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,26 +71,36 @@ func.func @dont_fold_reshape_with_not_full_load() {
// -----

#pipeline_layout = #hal.pipeline.layout<constants = 3, bindings = [
#hal.pipeline.binding<storage_buffer>,
#hal.pipeline.binding<storage_buffer>
]>
// CHECK-LABEL: func.func @dont_fold_dynamic_reshape()
func.func @dont_fold_dynamic_reshape() {
func.func @fold_dynamic_reshape() {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim0 = hal.interface.constant.load layout(#pipeline_layout) ordinal(0) : index
%dim1 = hal.interface.constant.load layout(#pipeline_layout) ordinal(1) : index
%dim2 = hal.interface.constant.load layout(#pipeline_layout) ordinal(2) : index
%1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) : !flow.dispatch.tensor<readonly:tensor<?x?x96xf32>>{%dim0, %dim1}
%2 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) : !flow.dispatch.tensor<writeonly:tensor<?x12x8xf32>>{%dim2}
%2 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) : !flow.dispatch.tensor<writeonly:tensor<?x12x8xf32>>{%dim2}
%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>
// CHECK: tensor.collapse_shape
// CHECK: tensor.expand_shape
%4 = tensor.collapse_shape %3 [[0, 1], [2]] : tensor<?x?x96xf32> into tensor<?x96xf32>
%dyn = tensor.dim %4, %c0 : tensor<?x96xf32>
%5 = tensor.expand_shape %4 [[0], [1, 2]] output_shape [%dyn, 12, 8] : tensor<?x96xf32> into tensor<?x12x8xf32>
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}
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}
return
}
// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 * s1)>
// CHECK: func.func @fold_dynamic_reshape()
// CHECK-DAG: %[[CST0:.+]] = hal.interface.constant.load layout(#{{.+}}) ordinal(0)
// CHECK-DAG: %[[CST1:.+]] = hal.interface.constant.load layout(#{{.+}}) ordinal(1)
// CHECK-DAG: %[[CST2:.+]] = hal.interface.constant.load layout(#{{.+}}) ordinal(2)
// CHECK: %[[COLLAPSED:.+]] = affine.apply #[[MAP]]()[%[[CST0]], %[[CST1]]]
// CHECK: %[[IN_BINDING:.+]] = hal.interface.binding.subspan
// CHECK-SAME: binding(0) : !flow.dispatch.tensor<readonly:tensor<?x96xf32>>{%[[COLLAPSED]]}
// CHECK: %[[OUT_BINDING:.+]] = hal.interface.binding.subspan
// CHECK-SAME: binding(1) : !flow.dispatch.tensor<writeonly:tensor<?x96xf32>>{%[[CST2]]}
// CHECK: %[[IN:.+]] = flow.dispatch.tensor.load %[[IN_BINDING]]
// CHECK: flow.dispatch.tensor.store %[[IN]], %[[OUT_BINDING]]

// -----

Expand Down
74 changes: 73 additions & 1 deletion compiler/src/iree/compiler/Codegen/Common/test/propagate_reshapes_by_expansion.mlir
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-propagate-reshapes-by-expansion))" --split-input-file %s | FileCheck %s
// 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

func.func @reshape_and_lowering_config(%src: tensor<3x4xf16>, %dest: tensor<12xf16>, %dest2: tensor<12xf16>) -> tensor<12xf16> {
%collapse = tensor.collapse_shape %src [[0, 1]] : tensor<3x4xf16> into tensor<12xf16>
Expand All @@ -14,3 +14,75 @@ func.func @reshape_and_lowering_config(%src: tensor<3x4xf16>, %dest: tensor<12xf
// CHECK: linalg.copy
// CHECK-SAME: lowering_config = #iree_gpu.derived_thread_config
// CHECK-SAME: ins(%[[COLLAPSE]]

// -----

#pipeline_layout = #hal.pipeline.layout<constants = 1, bindings = [
#hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">], flags = Indirect>
func.func @fold_collapse_into_loads_dynamic() -> tensor<?x32xf32> {
%c0 = arith.constant 0 : index
%0 = hal.interface.constant.load layout(#pipeline_layout) ordinal(0) : index
%1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
flags("ReadOnly|Indirect") : !flow.dispatch.tensor<readonly:tensor<2x?x32xf32>>{%0}
%2 = flow.dispatch.tensor.load %1, offsets = [0, 0, 0], sizes = [2, %0, 32], strides = [1, 1, 1]
: !flow.dispatch.tensor<readonly:tensor<2x?x32xf32>>{%0} -> tensor<2x?x32xf32>
%3 = tensor.collapse_shape %2 [[0, 1], [2]] : tensor<2x?x32xf32> into tensor<?x32xf32>
return %3 : tensor<?x32xf32>
}
// CHECK-LABEL: func @fold_collapse_into_loads_dynamic()
// CHECK: %[[CONST:.+]] = hal.interface.constant.load
// CHECK: %[[SHAPE:.+]] = affine.apply affine_map<()[s0] -> (s0 * 2)>()[%[[CONST]]]
// CHECK: %[[SUBSPAN:.+]] = hal.interface.binding.subspan
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x32xf32>>{%[[SHAPE]]}
// CHECK: %[[LOAD:.+]] = flow.dispatch.tensor.load %[[SUBSPAN]]
// CHECK-SAME: offsets = [0, 0], sizes = [%[[SHAPE]], 32], strides = [1, 1]
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x32xf32>>{%[[SHAPE]]}

// -----

#pipeline_layout = #hal.pipeline.layout<constants = 2, bindings = [
#hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">], flags = Indirect>
func.func @fold_expand_into_loads_dynamic() -> tensor<2x?x16x32xf32> {
%c0 = arith.constant 0 : index
%0 = hal.interface.constant.load layout(#pipeline_layout) ordinal(0) : index
%1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
flags("ReadOnly|Indirect") : !flow.dispatch.tensor<readonly:tensor<2x?x32xf32>>{%0}
%2 = flow.dispatch.tensor.load %1, offsets = [0, 0, 0], sizes = [2, %0, 32], strides = [1, 1, 1]
: !flow.dispatch.tensor<readonly:tensor<2x?x32xf32>>{%0} -> tensor<2x?x32xf32>
%3 = affine.apply affine_map<()[s0] -> (s0 floordiv 2)>()[%0]
%4 = tensor.expand_shape %2 [[0], [1, 2], [3]] output_shape [2, %3, 16, 32] : tensor<2x?x32xf32> into tensor<2x?x16x32xf32>
return %4 : tensor<2x?x16x32xf32>
}
// CHECK-LABEL: func @fold_expand_into_loads_dynamic()
// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
// CHECK-DAG: %[[CONST:.+]] = hal.interface.constant.load
// CHECK: %[[SHAPE:.+]] = arith.divui %[[CONST]], %[[C16]]
// CHECK: %[[SUBSPAN:.+]] = hal.interface.binding.subspan
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<2x?x16x32xf32>>{%[[SHAPE]]}
// CHECK: %[[LOAD:.+]] = flow.dispatch.tensor.load %[[SUBSPAN]]
// CHECK-SAME: offsets = [0, 0, 0, 0], sizes = [2, %[[SHAPE]], 16, 32], strides = [1, 1, 1, 1]
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<2x?x16x32xf32>>{%[[SHAPE]]}

// -----

#pipeline_layout = #hal.pipeline.layout<constants = 1, bindings = [
#hal.pipeline.binding<storage_buffer, Indirect>], flags = Indirect>
func.func @fold_collapse_into_stores_dynamic(%arg0 : tensor<2x?x32xf32>) {
%c0 = arith.constant 0 : index
%0 = hal.interface.constant.load layout(#pipeline_layout) ordinal(0) : index
%1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
flags("ReadOnly|Indirect") : !flow.dispatch.tensor<writeonly:tensor<?x32xf32>>{%0}
%2 = tensor.collapse_shape %arg0 [[0, 1], [2]] : tensor<2x?x32xf32> into tensor<?x32xf32>
flow.dispatch.tensor.store %2, %1, offsets = [0, 0], sizes = [%0, 32], strides = [1, 1]
: tensor<?x32xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x32xf32>>{%0}
return
}
// CHECK-LABEL: func @fold_collapse_into_stores_dynamic(
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK: %[[CONST:.+]] = hal.interface.constant.load
// CHECK: %[[SHAPE:.+]] = arith.divui %[[CONST]], %[[C2]]
// CHECK: %[[SUBSPAN:.+]] = hal.interface.binding.subspan
// CHECK-SAME: !flow.dispatch.tensor<writeonly:tensor<2x?x32xf32>>{%[[SHAPE]]}
// CHECK: flow.dispatch.tensor.store %{{.+}}, %[[SUBSPAN]]
// CHECK-SAME: offsets = [0, 0, 0], sizes = [2, %[[SHAPE]], 32], strides = [1, 1, 1]
// CHECK-SAME: !flow.dispatch.tensor<writeonly:tensor<2x?x32xf32>>{%[[SHAPE]]}
Loading

0 comments on commit 00104b5

Please sign in to comment.