[DispatchCreation] Allow hoisting of set_encoding through generics with a transposed output map (#24800)
This patch enables more opportunities to hoist `set_encoding` through
`linalg.generic` with non-identity output maps, e.g. cases when dequant
and transpose are fused together.
Assisted-by: Claude Code
Signed-off-by: Alex Denisov <denisov@roofline.ai>
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_x86_64.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_x86_64.mlir
index 4ccb59b..2979239 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_x86_64.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_x86_64.mlir
@@ -1560,6 +1560,54 @@
// -----
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
+#map4 = affine_map<(d0, d1, d2) -> (d1, d0)>
+#encoding_ibmap = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, [#map1, #map3], #map2]>
+#encoding_bmap = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, [#map1, #map4], #map2]>
+#encoding = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2]>
+func.func @dequantization_transposed(
+ %7: tensor<2x11008x128xi8, #encoding_ibmap>,
+ %8: tensor<2x11008xf32, #encoding_bmap>,
+ %9: tensor<2x11008xf32, #encoding_bmap>
+) -> tensor<11008x2x128xf32, #encoding> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="x86_64-xyz-xyz", cpu_features="+avx512f", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %13 = tensor.empty() : tensor<11008x2x128xf32, #encoding>
+ %14 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d1, d0, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%7, %8, %9 : tensor<2x11008x128xi8, #encoding_ibmap>, tensor<2x11008xf32, #encoding_bmap>, tensor<2x11008xf32, #encoding_bmap>) outs(%13 : tensor<11008x2x128xf32, #encoding>) {
+ ^bb0(%in: i8, %in_0: f32, %in_1: f32, %out: f32):
+ %21 = arith.extui %in : i8 to i32
+ %22 = arith.uitofp %21 : i32 to f32
+ %23 = arith.subf %22, %in_1 : f32
+ %24 = arith.mulf %23, %in_0 : f32
+ linalg.yield %24 : f32
+ } -> tensor<11008x2x128xf32, #encoding>
+ return %14 : tensor<11008x2x128xf32, #encoding>
+}
+
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d3)>
+// CHECK-LABEL: func.func @dequantization_transposed(
+// CHECK-SAME: %[[WEIGHT:[a-zA-Z0-9]+]]: tensor<11008x1x128x16x1xi8>
+// CHECK-SAME: %[[WEIGHT_SCALES:[a-zA-Z0-9]+]]: tensor<11008x1x16xf32>
+// CHECK-SAME: %[[WEIGHT_ZPS:[a-zA-Z0-9]+]]: tensor<11008x1x16xf32>
+// CHECK-SAME: ) -> tensor<11008x1x128x16x1xf32>
+// CHECK-DAG: %[[EMPTY_WEIGHT:.+]] = tensor.empty() : tensor<11008x1x128x16x1xf32>
+// CHECK-DAG: %[[WEIGHT_DEQUANT:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP1]], #[[$MAP]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]
+// CHECK-SAME: ins(%[[WEIGHT]], %[[WEIGHT_SCALES]], %[[WEIGHT_ZPS]] : tensor<11008x1x128x16x1xi8>, tensor<11008x1x16xf32>, tensor<11008x1x16xf32>)
+// CHECK-SAME: outs(%[[EMPTY_WEIGHT]] : tensor<11008x1x128x16x1xf32>)
+// CHECK: arith.extui
+// CHECK: arith.uitofp
+// CHECK: arith.subf
+// CHECK: arith.mulf
+// CHECK: return %[[WEIGHT_DEQUANT]]
+
+// -----
+
#encoding = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [[affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>], iteration_sizes = [2, 128, 64, ?]>
#encoding_bcast = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [[affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>, affine_map<(d0, d1, d2) -> (d1, d2)>], affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>], iteration_sizes = [2, 128, 64, ?]>
func.func @broadcast_batch(
diff --git a/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp b/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp
index 4796602..9c3b674 100644
--- a/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp
+++ b/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp
@@ -27,6 +27,7 @@
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
+#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
@@ -49,7 +50,7 @@
/// 1. Have a single result.
/// 2. Have single use.
/// 3. Have all parallel iterators.
-/// 4. Have an identity output indexing map.
+/// 4. Have a permutation (which includes identity) output indexing map.
/// 5. Have a tensor.empty init operand.
/// 6. Have as many indexing map dims as there are results in the encoding's
/// bcast_map.
@@ -58,6 +59,13 @@
/// `genericOp`, and replaces the op with an encoded version. If any of
/// the above conditions are false, then it returns failure.
///
+/// When the output map is a non-identity permutation (e.g. a transpose fused
+/// into the generic's output indexing map by elementwise fusion), each
+/// operand's indexing map is re-expressed in terms of the *output*'s index
+/// space, rather than the generic's raw loop space, by composing with
+/// the inverse of the output permutation, before it is appended to the
+/// encoding's map chain.
+///
/// Note: The bcast_map on the set_encoding op must be identity or absent.
/// The implementation should work for cases where it is not, but it is
/// unexpected in IREE compilation to find such cases, and it will not
@@ -84,9 +92,11 @@
}
AffineMap outputMap =
genericOp.getMatchingIndexingMap(genericOp.getDpsInitOperand(0));
- if (!outputMap.isIdentity()) {
- return rewriter.notifyMatchFailure(genericOp, "output map not identity");
+ if (!outputMap.isPermutation()) {
+ return rewriter.notifyMatchFailure(genericOp,
+ "output map not a permutation");
}
+ AffineMap invOutputMap = inversePermutation(outputMap);
RankedTensorType encodedType = encodingOp.getResultType();
auto encoding = cast<IREE::Encoding::EncodingAttr>(encodedType.getEncoding());
@@ -116,12 +126,15 @@
Location loc = genericOp->getLoc();
SmallVector<Value> encodedOperands;
for (OpOperand *operand : genericOp.getDpsInputOperands()) {
- // Append the operand's indexing map to the encoding's user indexing maps.
+ // Append the operand's indexing map to the encoding's user indexing
+ // maps, re-expressed in terms of the output's index space by
+ // composing with the inverse of the output permutation.
AffineMap operandMap = genericOp.getMatchingIndexingMap(operand);
+ AffineMap remappedOperandMap = operandMap.compose(invOutputMap);
// Create new encoding and set encoding on the operand.
IREE::Encoding::EncodingAttr newEncoding =
- encoding.cloneWithNewOperandIndexingMap(operandMap);
+ encoding.cloneWithNewOperandIndexingMap(remappedOperandMap);
auto operandType = cast<RankedTensorType>(operand->get().getType());
auto resType = RankedTensorType::get(
operandType.getShape(), operandType.getElementType(), newEncoding);
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir b/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
index 95a1a7a..88d2381 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
@@ -683,18 +683,18 @@
// can be rematerialized as tensor.dim on the inputs (which fold to constants
// for static dimensions).
-#map_remat = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-#map_remat1 = affine_map<(d0, d1, d2) -> (d0, d1)>
-#map_remat2 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
-#map_remat3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
-#map_remat4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
-#encoding_remat = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map_remat2, #map_remat3, #map_remat4], iteration_sizes = [2, ?, ?, ?]>
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+#encoding = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map2, #map3, #map4], iteration_sizes = [2, ?, ?, ?]>
util.func public @bubble_with_rematerialized_encoding_dims(
- %arg0: tensor<2x11008x128xi8>, %arg1: tensor<2x11008xf32>, %arg2: tensor<2x11008xf32>) -> tensor<2x11008x128xf32, #encoding_remat> {
- %6 = flow.dispatch.region -> (tensor<2x11008x128xf32, #encoding_remat>) {
+ %arg0: tensor<2x11008x128xi8>, %arg1: tensor<2x11008xf32>, %arg2: tensor<2x11008xf32>) -> tensor<2x11008x128xf32, #encoding> {
+ %6 = flow.dispatch.region -> (tensor<2x11008x128xf32, #encoding>) {
%8 = tensor.empty() : tensor<2x11008x128xf32>
%11 = linalg.generic
- {indexing_maps = [#map_remat, #map_remat1, #map_remat1, #map_remat],
+ {indexing_maps = [#map, #map1, #map1, #map],
iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0, %arg1, %arg2 : tensor<2x11008x128xi8>, tensor<2x11008xf32>, tensor<2x11008xf32>)
outs(%8 : tensor<2x11008x128xf32>) {
@@ -713,10 +713,10 @@
%m = tensor.dim %11, %c0 : tensor<2x11008x128xf32>
%n = tensor.dim %11, %c1 : tensor<2x11008x128xf32>
%k = tensor.dim %11, %c2 : tensor<2x11008x128xf32>
- %13 = iree_encoding.set_encoding %11 encoding_dims{%m, %n, %k} : tensor<2x11008x128xf32> -> tensor<2x11008x128xf32, #encoding_remat>
- flow.return %13 : tensor<2x11008x128xf32, #encoding_remat>
+ %13 = iree_encoding.set_encoding %11 encoding_dims{%m, %n, %k} : tensor<2x11008x128xf32> -> tensor<2x11008x128xf32, #encoding>
+ flow.return %13 : tensor<2x11008x128xf32, #encoding>
}
- util.return %6 : tensor<2x11008x128xf32, #encoding_remat>
+ util.return %6 : tensor<2x11008x128xf32, #encoding>
}
// The tensor.dim ops are rematerialized from the first input, folding to constants.
@@ -921,3 +921,197 @@
// CHECK: %[[BCAST:.+]] = linalg.generic
// CHECK-SAME: ins(%[[UNSET]]
// CHECK: flow.return %[[BCAST]]
+
+
+// -----
+
+// Verify bubbling of a set_encoding through a dequant generic whose output
+// indexing map is a non-identity permutation (e.g. a transpose fused into the
+// producer by elementwise fusion).
+// Each operand's own indexing map should be composed with the inverse of the
+// output permutation before being appended to the encoding's map chain.
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
+#map3 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map5 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+#encoding = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map3, #map4, #map5]>
+util.func public @bubble_through_dequant_and_transpose(
+ %arg0: tensor<2x11008x128xi8>, %arg1: tensor<2x11008xf32>, %arg2: tensor<2x11008xf32>) -> tensor<11008x2x128xf32, #encoding> {
+ %6 = flow.dispatch.region -> (tensor<11008x2x128xf32, #encoding>) {
+ %8 = tensor.empty() : tensor<11008x2x128xf32>
+ %11 = linalg.generic
+ {indexing_maps = [#map, #map1, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "parallel"]}
+ ins(%arg0, %arg1, %arg2 : tensor<2x11008x128xi8>, tensor<2x11008xf32>, tensor<2x11008xf32>)
+ outs(%8 : tensor<11008x2x128xf32>) {
+ ^bb0(%in: i8, %in_0: f32, %in_1: f32, %out: f32):
+ %18 = arith.extui %in : i8 to i32
+ %19 = arith.uitofp %18 : i32 to f32
+ %20 = arith.subf %19, %in_1 : f32
+ %21 = arith.mulf %20, %in_0 : f32
+ linalg.yield %21 : f32
+ } -> tensor<11008x2x128xf32>
+ %13 = iree_encoding.set_encoding %11 : tensor<11008x2x128xf32> -> tensor<11008x2x128xf32, #encoding>
+ flow.return %13 : tensor<11008x2x128xf32, #encoding>
+ }
+ util.return %6 : tensor<11008x2x128xf32, #encoding>
+}
+
+// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
+// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2) -> (d1, d0)>
+// CHECK-DAG: #{{.*}} = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #{{.*}} = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP]], #[[MAP1]], #[[MAP2]]]>
+// CHECK-DAG: #[[$ENCODING_IBMAP:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP]], [#[[MAP1]], #[[MAP3]]], #[[MAP2]]]>
+// CHECK-DAG: #[[$ENCODING_BMAP:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#[[MAP]], [#[[MAP1]], #[[MAP4]]], #[[MAP2]]]>
+// CHECK-LABEL: @bubble_through_dequant_and_transpose
+// CHECK-SAME: %[[ARG0:.+]]: tensor<2x11008x128xi8>,
+// CHECK-SAME: %[[ARG1:.+]]: tensor<2x11008xf32>, %[[ARG2:.+]]: tensor<2x11008xf32>
+// CHECK-DAG: %[[SET_ENCODING0:.+]] = iree_encoding.set_encoding %[[ARG0]] : tensor<2x11008x128xi8> -> tensor<2x11008x128xi8, #[[$ENCODING_IBMAP]]>
+// CHECK-DAG: %[[SET_ENCODING1:.+]] = iree_encoding.set_encoding %[[ARG1]] : tensor<2x11008xf32> -> tensor<2x11008xf32, #[[$ENCODING_BMAP]]>
+// CHECK-DAG: %[[SET_ENCODING2:.+]] = iree_encoding.set_encoding %[[ARG2]] : tensor<2x11008xf32> -> tensor<2x11008xf32, #[[$ENCODING_BMAP]]>
+// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region
+// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<11008x2x128xf32, #[[$ENCODING]]>
+// CHECK: %[[DEQUANT:.+]] = linalg.generic {{.*}} ins(%[[SET_ENCODING0]], %[[SET_ENCODING1]], %[[SET_ENCODING2]] : {{.*}} outs(%[[INIT]] :
+// CHECK: flow.return %[[DEQUANT]]
+// CHECK: }
+// CHECK: util.return %[[DISPATCH]]
+
+// -----
+
+// Same bubbling as above, but with an output permutation that is not its
+// own inverse (d0, d1, d2 -> d1, d2, d0).
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#encoding = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32], user_indexing_maps = [#map2, #map3, #map2]>
+util.func public @bubble_through_non_involutive_permutation(
+ %arg0: tensor<2x11008x128xf16>) -> tensor<11008x128x2xf32, #encoding> {
+ %6 = flow.dispatch.region -> (tensor<11008x128x2xf32, #encoding>) {
+ %8 = tensor.empty() : tensor<11008x128x2xf32>
+ %11 = linalg.generic
+ {indexing_maps = [#map, #map1],
+ iterator_types = ["parallel", "parallel", "parallel"]}
+ ins(%arg0 : tensor<2x11008x128xf16>)
+ outs(%8 : tensor<11008x128x2xf32>) {
+ ^bb0(%in: f16, %out: f32):
+ %18 = arith.extf %in : f16 to f32
+ linalg.yield %18 : f32
+ } -> tensor<11008x128x2xf32>
+ %13 = iree_encoding.set_encoding %11 : tensor<11008x128x2xf32> -> tensor<11008x128x2xf32, #encoding>
+ flow.return %13 : tensor<11008x128x2xf32, #encoding>
+ }
+ util.return %6 : tensor<11008x128x2xf32, #encoding>
+}
+
+// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
+// CHECK-DAG: #{{.*}} = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #{{.*}} = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
+// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32], user_indexing_maps = [#[[MAP]], #[[MAP1]], #[[MAP]]]>
+// CHECK-DAG: #[[$ENCODING_IBMAP:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32], user_indexing_maps = [#[[MAP]], [#[[MAP1]], #[[MAP2]]], #[[MAP]]]>
+// CHECK-LABEL: @bubble_through_non_involutive_permutation
+// CHECK-SAME: %[[ARG0:.+]]: tensor<2x11008x128xf16>
+// CHECK: %[[SET_ENCODING0:.+]] = iree_encoding.set_encoding %[[ARG0]] : tensor<2x11008x128xf16> -> tensor<2x11008x128xf16, #[[$ENCODING_IBMAP]]>
+// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region
+// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<11008x128x2xf32, #[[$ENCODING]]>
+// CHECK: %[[DEQUANT:.+]] = linalg.generic {{.*}} ins(%[[SET_ENCODING0]] : {{.*}} outs(%[[INIT]] :
+// CHECK: flow.return %[[DEQUANT]]
+// CHECK: }
+// CHECK: util.return %[[DISPATCH]]
+
+// -----
+
+// A reduction output map isn't a permutation, so the set_encoding must not bubble through it.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d1)>
+#map2 = affine_map<(d0, d1) -> (d0)>
+#encoding = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2]>
+util.func public @dont_bubble_through_reduction_output_map(
+ %arg0: tensor<8x4096xf32>) -> tensor<8xf32, #encoding> {
+ %0 = flow.dispatch.region -> (tensor<8xf32, #encoding>) {
+ %2 = tensor.empty() : tensor<8xf32>
+ %3 = linalg.generic {indexing_maps = [#map, #map2], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<8x4096xf32>) outs(%2 : tensor<8xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %4 = arith.addf %in, %out : f32
+ linalg.yield %4 : f32
+ } -> tensor<8xf32>
+ %5 = iree_encoding.set_encoding %3 : tensor<8xf32> -> tensor<8xf32, #encoding>
+ flow.return %5 : tensor<8xf32, #encoding>
+ }
+ util.return %0 : tensor<8xf32, #encoding>
+}
+// CHECK-LABEL: @dont_bubble_through_reduction_output_map
+// CHECK: flow.dispatch.region
+// CHECK: tensor.empty
+// CHECK: linalg.generic
+// CHECK: iree_encoding.set_encoding
+// CHECK: flow.return
+// CHECK: util.return
+
+// -----
+
+// encoding_dims rematerialization must still work when the raw inputs' dim order differs from the transposed generic output's.
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
+#map3 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
+#map4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map5 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+#encoding = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map3, #map4, #map5], iteration_sizes = [2, ?, ?, ?]>
+util.func public @bubble_with_rematerialized_encoding_dims_transposed(
+ %arg0: tensor<2x11008x128xi8>, %arg1: tensor<2x11008xf32>, %arg2: tensor<2x11008xf32>) -> tensor<11008x2x128xf32, #encoding> {
+ %6 = flow.dispatch.region -> (tensor<11008x2x128xf32, #encoding>) {
+ %8 = tensor.empty() : tensor<11008x2x128xf32>
+ %11 = linalg.generic
+ {indexing_maps = [#map, #map1, #map1, #map2],
+ iterator_types = ["parallel", "parallel", "parallel"]}
+ ins(%arg0, %arg1, %arg2 : tensor<2x11008x128xi8>, tensor<2x11008xf32>, tensor<2x11008xf32>)
+ outs(%8 : tensor<11008x2x128xf32>) {
+ ^bb0(%in: i8, %in_0: f32, %in_1: f32, %out: f32):
+ %18 = arith.extui %in : i8 to i32
+ %19 = arith.uitofp %18 : i32 to f32
+ %20 = arith.subf %19, %in_1 : f32
+ %21 = arith.mulf %20, %in_0 : f32
+ linalg.yield %21 : f32
+ } -> tensor<11008x2x128xf32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %m = tensor.dim %11, %c0 : tensor<11008x2x128xf32>
+ %n = tensor.dim %11, %c1 : tensor<11008x2x128xf32>
+ %k = tensor.dim %11, %c2 : tensor<11008x2x128xf32>
+ %13 = iree_encoding.set_encoding %11 encoding_dims{%m, %n, %k} : tensor<11008x2x128xf32> -> tensor<11008x2x128xf32, #encoding>
+ flow.return %13 : tensor<11008x2x128xf32, #encoding>
+ }
+ util.return %6 : tensor<11008x2x128xf32, #encoding>
+}
+
+// The tensor.dim ops rematerialize from the raw i8 input and fold to constants, keeping the transposed output's (N, batch, K) value order.
+// CHECK-DAG: #[[$ENCODING:.+]] = #iree_encoding.encoding<operand_index = 1 : index, op_type = matmul, element_types = [f32, f32, f32]
+// CHECK-LABEL: @bubble_with_rematerialized_encoding_dims_transposed
+// CHECK-SAME: %[[ARG0:.+]]: tensor<2x11008x128xi8>,
+// CHECK-SAME: %[[ARG1:.+]]: tensor<2x11008xf32>, %[[ARG2:.+]]: tensor<2x11008xf32>
+// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[C11008:.+]] = arith.constant 11008 : index
+// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK-DAG: %[[SET_ENCODING0:.+]] = iree_encoding.set_encoding %[[ARG0]] encoding_dims{%[[C11008]], %[[C2]], %[[C128]]}
+// CHECK-DAG: %[[SET_ENCODING1:.+]] = iree_encoding.set_encoding %[[ARG1]] encoding_dims{%[[C11008]], %[[C2]], %[[C128]]}
+// CHECK-DAG: %[[SET_ENCODING2:.+]] = iree_encoding.set_encoding %[[ARG2]] encoding_dims{%[[C11008]], %[[C2]], %[[C128]]}
+// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region
+// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<11008x2x128xf32, #[[$ENCODING]]>
+// CHECK: %[[DEQUANT:.+]] = linalg.generic {{.*}} ins(%[[SET_ENCODING0]], %[[SET_ENCODING1]], %[[SET_ENCODING2]] : {{.*}} outs(%[[INIT]] :
+// CHECK: flow.return %[[DEQUANT]]
+// CHECK: }
+// CHECK: util.return %[[DISPATCH]]
