[DispatchCreation] Propagate tensor encodings outside dispatch regions (#24966) Encoding propagation is currently tied to dispatch hoisting, which prevents layout assignment from using the same rewrite rules on the surrounding tensor graph. Separate where propagation may happen from the rewrites themselves, and add the function pass `iree-dispatch-creation-propagate-data-tiling-encodings` for propagation outside dispatches. Hoisting keeps its existing behavior. The new pass stops at existing dispatches and block boundaries, and constant CSE stays disabled because Flow regions do not materialize constants themselves. - `iree-dispatch-creation-propagate-encodings` runs in the dispatch-time data-tiling pipeline and only swaps `set_encoding` with `tensor.cast` producers, so adding these rewrites to it would change that pipeline. They stay separate until bubbling through generic ops goes through the encoding propagation interfaces (TODO(#20179)). - Bubbling now stops at `set_encoding` ops whose encoding is not an `EncodingAttr`, such as `#iree_encoding.padding`, instead of crashing. `iree-dispatch-creation-hoist-encoding-ops` could already hit this within dispatches; both passes have a regression test. No pipeline runs the new pass yet. #24969 adds it to the experimental early data-tiling route. This is one of five PRs adding an experimental CPU path that materializes data-tiling layouts before dispatch creation, so that `linalg.mmt4d` and data-tiled convolutions fuse with their elementwise epilogues and result unpacks. Assisted by AI.
diff --git a/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp b/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp index 81312b9..9ab1144 100644 --- a/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp +++ b/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp
@@ -43,6 +43,7 @@ namespace mlir::iree_compiler::DispatchCreation { #define GEN_PASS_DEF_HOISTENCODINGOPSPASS +#define GEN_PASS_DEF_PROPAGATEENCODINGSTHROUGHGENERICOPSPASS #include "iree/compiler/DispatchCreation/Passes.h.inc" /// Bubbles a SetEncodingOp up through a linalg::GenericOp. The `genericOp` @@ -54,6 +55,7 @@ /// 5. Have a tensor.empty init operand. /// 6. Have as many indexing map dims as there are results in the encoding's /// bcast_map. +/// The encoding of `encodingOp` must be an `EncodingAttr`. /// /// This function creates SetEncoding ops on all of the inputs to the /// `genericOp`, and replaces the op with an encoded version. If any of @@ -99,7 +101,11 @@ AffineMap invOutputMap = inversePermutation(outputMap); RankedTensorType encodedType = encodingOp.getResultType(); - auto encoding = cast<IREE::Encoding::EncodingAttr>(encodedType.getEncoding()); + auto encoding = + dyn_cast<IREE::Encoding::EncodingAttr>(encodedType.getEncoding()); + if (!encoding) { + return rewriter.notifyMatchFailure(encodingOp, "expected an EncodingAttr"); + } AffineMap lastMap = encoding.getLastMapForOperandIndex(); if (outputMap.getNumDims() != lastMap.getNumResults()) { return rewriter.notifyMatchFailure( @@ -193,43 +199,83 @@ } namespace { +/// Where encoding propagation is allowed to move encoding ops. +enum class EncodingPropagationScope { WithinDispatch, OutsideDispatch }; +} // namespace + +/// Returns true if an encoding may be propagated between `producer` and +/// `consumer` under `scope`. Within dispatches, both operations must be inside +/// a dispatch. Outside dispatches, both must be outside every dispatch and in +/// the same block, so existing dispatches and region boundaries remain encoding +/// boundaries. +static bool canPropagateBetween(Operation *producer, Operation *consumer, + EncodingPropagationScope scope) { + bool producerOutside = IREE::Flow::isNonNullAndOutsideDispatch(producer); + bool consumerOutside = IREE::Flow::isNonNullAndOutsideDispatch(consumer); + if (scope == EncodingPropagationScope::WithinDispatch) { + return !producerOutside && !consumerOutside; + } + return producerOutside && consumerOutside && + producer->getBlock() == consumer->getBlock(); +} + +namespace { /// Pass declaration. -struct HoistEncodingOpsPass +struct HoistEncodingOpsPass final : impl::HoistEncodingOpsPassBase<HoistEncodingOpsPass> { using Base::Base; void runOnOperation() override; }; -/// Pattern to bubble SetEncoding ops upwards through producers. This pattern -/// runs until bubbling is not possible, or until the SetEncoding op is outside -/// of a dispatch. +// TODO(#20179): Merge with PropagateEncodingsPass once bubbling through +// generic ops goes through the encoding propagation interfaces. +struct PropagateEncodingsThroughGenericOpsPass final + : impl::PropagateEncodingsThroughGenericOpsPassBase< + PropagateEncodingsThroughGenericOpsPass> { + void runOnOperation() override; +}; + +/// Pattern to bubble SetEncoding ops upwards through producers. Within +/// dispatches, it only bubbles through producers in the same dispatch region, +/// which prepares the SetEncoding op for hoisting. Outside dispatches, it only +/// bubbles through producers that are outside every dispatch and in the same +/// block. struct BubbleUpSetEncodingOp : OpRewritePattern<IREE::Encoding::SetEncodingOp> { - using Base::Base; + BubbleUpSetEncodingOp(MLIRContext *context, EncodingPropagationScope scope) + : Base(context), scope(scope) {} LogicalResult matchAndRewrite(IREE::Encoding::SetEncodingOp encodingOp, PatternRewriter &rewriter) const override { - if (IREE::Flow::isNonNullAndOutsideDispatch(encodingOp)) { - return failure(); - } - // Fail if the encodingOp is not in the same dispatch as its producer. Operation *producer = encodingOp.getSource().getDefiningOp(); if (!producer) { - return failure(); + return rewriter.notifyMatchFailure(encodingOp, "source has no producer"); } - auto dispatch = producer->getParentOfType<IREE::Flow::DispatchRegionOp>(); - if (!dispatch || - dispatch != - encodingOp->getParentOfType<IREE::Flow::DispatchRegionOp>()) { - return failure(); + if (!canPropagateBetween(producer, encodingOp, scope)) { + return rewriter.notifyMatchFailure( + encodingOp, "producer is outside the propagation scope"); + } + if (scope == EncodingPropagationScope::WithinDispatch) { + // Only bubble through producers in the same dispatch region. + auto dispatch = producer->getParentOfType<IREE::Flow::DispatchRegionOp>(); + if (!dispatch || + dispatch != + encodingOp->getParentOfType<IREE::Flow::DispatchRegionOp>()) { + return rewriter.notifyMatchFailure( + encodingOp, "producer is not in the same dispatch region"); + } } return bubbleUpSetEncoding(rewriter, encodingOp->getOpOperand(0)); } + +private: + EncodingPropagationScope scope; }; /// Pattern to sink UnsetEncoding ops down through consumers. struct SinkUnsetEncodingOp : OpRewritePattern<IREE::Encoding::UnsetEncodingOp> { - using Base::Base; + SinkUnsetEncodingOp(MLIRContext *context, EncodingPropagationScope scope) + : Base(context), scope(scope) {} LogicalResult matchAndRewrite(IREE::Encoding::UnsetEncodingOp encodingOp, PatternRewriter &rewriter) const override { @@ -238,10 +284,9 @@ } OpOperand *consumerOperand = &(*encodingOp->getUses().begin()); Operation *consumer = consumerOperand->getOwner(); - if (IREE::Flow::isNonNullAndOutsideDispatch(encodingOp) || - IREE::Flow::isNonNullAndOutsideDispatch(consumer)) { + if (!canPropagateBetween(encodingOp, consumer, scope)) { return rewriter.notifyMatchFailure( - encodingOp, "expected that both operations are inside dispatch"); + encodingOp, "operations are outside the propagation scope"); } auto propagationAttrInterface = @@ -284,22 +329,40 @@ rewriter.replaceOp(consumer, maybeResult->replacements); return success(); } + +private: + EncodingPropagationScope scope; }; } // namespace -/// Create dispatch.region Ops based on a fusion heuristic. +/// Bubbles set_encoding and sinks unset_encoding ops in `root` within `scope`. +static LogicalResult propagateEncodings(Operation *root, + EncodingPropagationScope scope) { + MLIRContext *context = root->getContext(); + RewritePatternSet patterns(context); + patterns.insert<BubbleUpSetEncodingOp, SinkUnsetEncodingOp>(context, scope); + // Flow regions are not isolated from above and do not materialize constants + // themselves, so constant CSE would move constants out of existing dispatch + // regions and turn them into dispatch captures. + GreedyRewriteConfig config; + config.enableConstantCSE(false); + return applyPatternsGreedily(root, std::move(patterns), config); +} + +void PropagateEncodingsThroughGenericOpsPass::runOnOperation() { + if (failed(propagateEncodings(getOperation(), + EncodingPropagationScope::OutsideDispatch))) { + signalPassFailure(); + } +} + void HoistEncodingOpsPass::runOnOperation() { MLIRContext *ctx = &getContext(); ModuleOp moduleOp = getOperation(); - RewritePatternSet bubblingPatterns(ctx); - bubblingPatterns.insert<BubbleUpSetEncodingOp>(ctx); - bubblingPatterns.insert<SinkUnsetEncodingOp>(ctx); - GreedyRewriteConfig config; - config.enableConstantCSE(false); - if (failed(applyPatternsGreedily(moduleOp, std::move(bubblingPatterns), - config))) { + if (failed(propagateEncodings(moduleOp, + EncodingPropagationScope::WithinDispatch))) { return signalPassFailure(); } @@ -393,6 +456,8 @@ RewritePatternSet cleanPatterns(ctx); memref::populateResolveRankedShapedTypeResultDimsPatterns(cleanPatterns); IREE::Flow::DispatchRegionOp::getCanonicalizationPatterns(cleanPatterns, ctx); + GreedyRewriteConfig config; + config.enableConstantCSE(false); if (failed( applyPatternsGreedily(moduleOp, std::move(cleanPatterns), config))) { return signalPassFailure();
diff --git a/compiler/src/iree/compiler/DispatchCreation/Passes.td b/compiler/src/iree/compiler/DispatchCreation/Passes.td index 0b67715..8c23809 100644 --- a/compiler/src/iree/compiler/DispatchCreation/Passes.td +++ b/compiler/src/iree/compiler/DispatchCreation/Passes.td
@@ -424,6 +424,29 @@ ]; } +def PropagateEncodingsThroughGenericOpsPass : + InterfacePass<"iree-dispatch-creation-propagate-encodings-through-generic-ops", "mlir::FunctionOpInterface"> { + let summary = "Propagates data-tiling encodings through linalg.generic ops outside dispatches."; + let description = [{ + Bubbles set_encoding ops up through bit-extending and broadcasting + linalg.generic producers, and sinks unset_encoding ops down through + elementwise linalg.generic consumers, for operations that are outside + dispatches and in the same block. These are the same encoding rewrites that + iree-dispatch-creation-hoist-encoding-ops applies within dispatch regions. + Unsupported operations and encodings, multiple uses, existing dispatches and + region boundaries remain encoding boundaries. The pass neither hoists out of + dispatches nor propagates across calls. Unlike + iree-dispatch-creation-propagate-encodings, which only swaps set_encoding + ops with tensor.cast producers, it rewrites the generic ops around the + encoding ops. + }]; + let dependentDialects = [ + "mlir::linalg::LinalgDialect", + "mlir::tensor::TensorDialect", + "IREE::Encoding::IREEEncodingDialect", + ]; +} + def AnnotateDataTilingHintsPass : InterfacePass<"iree-dispatch-creation-annotate-data-tiling-hints", "mlir::FunctionOpInterface"> { let summary = "Adds data-tiling hint attribute to linalg operations.";
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/BUILD.bazel b/compiler/src/iree/compiler/DispatchCreation/test/BUILD.bazel index 7e19e92..e42d518 100644 --- a/compiler/src/iree/compiler/DispatchCreation/test/BUILD.bazel +++ b/compiler/src/iree/compiler/DispatchCreation/test/BUILD.bazel
@@ -58,6 +58,7 @@ "pipeline_tests_split_reduction.mlir", "propagate_bitcast.mlir", "propagate_encodings.mlir", + "propagate_encodings_through_generic_ops.mlir", "remove_tensor_barriers.mlir", "set_encoding.mlir", "set_encoding_padding.mlir",
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/CMakeLists.txt b/compiler/src/iree/compiler/DispatchCreation/test/CMakeLists.txt index e407742..08a75bc 100644 --- a/compiler/src/iree/compiler/DispatchCreation/test/CMakeLists.txt +++ b/compiler/src/iree/compiler/DispatchCreation/test/CMakeLists.txt
@@ -55,6 +55,7 @@ "pipeline_tests_split_reduction.mlir" "propagate_bitcast.mlir" "propagate_encodings.mlir" + "propagate_encodings_through_generic_ops.mlir" "remove_tensor_barriers.mlir" "set_encoding.mlir" "set_encoding_padding.mlir"
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 7751424..c0b4f9c 100644 --- a/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir +++ b/compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir
@@ -340,6 +340,36 @@ // ----- +// Do not bubble `set_encoding` operations with pad encodings through their +// producers. + +#map = affine_map<(d0, d1) -> (d0, d1)> +#encoding = #iree_encoding.padding<[0, 64]> +util.func public @dont_bubble_pad_encoding(%arg0: tensor<128x256xi8>) -> tensor<128x256xf32, #encoding> { + %0 = flow.dispatch.region -> (tensor<128x256xf32, #encoding>) { + %1 = tensor.empty() : tensor<128x256xf32> + %2 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} + ins(%arg0 : tensor<128x256xi8>) outs(%1 : tensor<128x256xf32>) { + ^bb0(%in: i8, %out: f32): + %3 = arith.extui %in : i8 to i32 + %4 = arith.uitofp %3 : i32 to f32 + linalg.yield %4 : f32 + } -> tensor<128x256xf32> + %5 = iree_encoding.set_encoding %2 : tensor<128x256xf32> -> tensor<128x256xf32, #encoding> + flow.return %5 : tensor<128x256xf32, #encoding> + } + util.return %0 : tensor<128x256xf32, #encoding> +} +// CHECK-LABEL: @dont_bubble_pad_encoding +// CHECK: %[[DISPATCH:.+]] = flow.dispatch.region +// CHECK: %[[DEQUANT:.+]] = linalg.generic +// CHECK-SAME: ins(%{{.+}} : tensor<128x256xi8>) +// CHECK: %[[SET_ENCODING:.+]] = iree_encoding.set_encoding %[[DEQUANT]] +// CHECK: flow.return %[[SET_ENCODING]] +// CHECK: return %[[DISPATCH]] + +// ----- + // Avoid hoisting `set_encoding` operations on scalar tensors. #encoding = #iree_encoding.testing<>
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/propagate_encodings_through_generic_ops.mlir b/compiler/src/iree/compiler/DispatchCreation/test/propagate_encodings_through_generic_ops.mlir new file mode 100644 index 0000000..f51fb14 --- /dev/null +++ b/compiler/src/iree/compiler/DispatchCreation/test/propagate_encodings_through_generic_ops.mlir
@@ -0,0 +1,307 @@ +// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(util.func(iree-dispatch-creation-propagate-encodings-through-generic-ops))" %s | FileCheck %s + +// The pass applies the rewrites of hoist-encoding-ops outside dispatches. +// hoist-encoding-ops itself only propagates within dispatches, so it must leave +// the encoding ops outside dispatches in place. +// RUN: iree-opt --split-input-file --iree-dispatch-creation-hoist-encoding-ops %s | FileCheck %s --check-prefix=HOIST + +#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]> +// CHECK-LABEL: @bubble_through_dequant( +// CHECK: %[[A:.*]] = iree_encoding.set_encoding %arg0 +// CHECK: %[[B:.*]] = iree_encoding.set_encoding %arg1 +// CHECK: %[[C:.*]] = iree_encoding.set_encoding %arg2 +// CHECK: linalg.generic {{.*}} ins(%[[A]], %[[B]], %[[C]] +// CHECK-SAME: tensor<2x11008x128xi8, # +// CHECK-NOT: iree_encoding.set_encoding +// CHECK: util.return +// HOIST-LABEL: @bubble_through_dequant( +// HOIST: %[[DEQUANT:.*]] = linalg.generic +// HOIST: iree_encoding.set_encoding %[[DEQUANT]] +util.func public @bubble_through_dequant( + %arg0: tensor<2x11008x128xi8>, %arg1: tensor<2x11008xf32>, %arg2: tensor<2x11008xf32>) -> tensor<2x11008x128xf32, #encoding> { + %8 = tensor.empty() : tensor<2x11008x128xf32> + %11 = linalg.generic + {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>) { + ^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<2x11008x128xf32> + %13 = iree_encoding.set_encoding %11 : tensor<2x11008x128xf32> -> tensor<2x11008x128xf32, #encoding> + util.return %13 : tensor<2x11008x128xf32, #encoding> +} + +// ----- + +#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d1, d2)> +#encoding = #iree_encoding.encoding<operand_index = 1 : 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, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]> +// CHECK-LABEL: @bubble_through_broadcast( +// CHECK: %[[A:.*]] = iree_encoding.set_encoding %arg0 : tensor<11008x128xf32> -> tensor<11008x128xf32, # +// CHECK: linalg.generic {{.*}} ins(%[[A]] +// CHECK-SAME: outs(%{{.*}} : tensor<2x11008x128xf32, # +// CHECK-NOT: iree_encoding.set_encoding +// CHECK: util.return +util.func public @bubble_through_broadcast( + %arg0: tensor<11008x128xf32>) -> tensor<2x11008x128xf32, #encoding> { + %8 = tensor.empty() : tensor<2x11008x128xf32> + %11 = linalg.generic + {indexing_maps = [#map1, #map], + iterator_types = ["parallel", "parallel", "parallel"]} + ins(%arg0 : tensor<11008x128xf32>) + outs(%8 : tensor<2x11008x128xf32>) { + ^bb0(%in: f32, %out: f32): + linalg.yield %in : f32 + } -> tensor<2x11008x128xf32> + %13 = iree_encoding.set_encoding %11 : tensor<2x11008x128xf32> -> tensor<2x11008x128xf32, #encoding> + util.return %13 : tensor<2x11008x128xf32, #encoding> +} + +// ----- + +// Padding encodings are not data-tiling encodings of a contraction, so they +// are not bubbled. +#map = affine_map<(d0, d1) -> (d0, d1)> +#encoding = #iree_encoding.padding<[0, 64]> +// CHECK-LABEL: @no_bubble_padding_encoding( +// CHECK: %[[DEQUANT:.+]] = linalg.generic +// CHECK-SAME: ins(%{{.+}} : tensor<128x256xi8>) +// CHECK: iree_encoding.set_encoding %[[DEQUANT]] +util.func public @no_bubble_padding_encoding(%arg0: tensor<128x256xi8>) -> tensor<128x256xf32, #encoding> { + %empty = tensor.empty() : tensor<128x256xf32> + %dequant = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} + ins(%arg0 : tensor<128x256xi8>) outs(%empty : tensor<128x256xf32>) { + ^bb0(%in: i8, %out: f32): + %ext = arith.extui %in : i8 to i32 + %fp = arith.uitofp %ext : i32 to f32 + linalg.yield %fp : f32 + } -> tensor<128x256xf32> + %encoded = iree_encoding.set_encoding %dequant : tensor<128x256xf32> -> tensor<128x256xf32, #encoding> + util.return %encoded : tensor<128x256xf32, #encoding> +} + +// ----- + +#map = affine_map<(d0, d1, d2) -> (d0, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d1, d2)> +#map2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#map3 = affine_map<(d0, d1) -> (d0, d1)> +#map4 = affine_map<(d0, d1) -> ()> +#encoding = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2]> +// CHECK-LABEL: @propagate_unset_encoding_through_generic_with_scalar( +// CHECK-NOT: iree_encoding.unset_encoding +// CHECK: %[[RESULT:.*]] = linalg.generic +// CHECK-SAME: ins(%arg0, %arg1 : tensor<4096x?xf32, #{{.*}}>, f32) +// CHECK: iree_encoding.unset_encoding %[[RESULT]] +// HOIST-LABEL: @propagate_unset_encoding_through_generic_with_scalar( +// HOIST: %[[RAW:.*]] = iree_encoding.unset_encoding +// HOIST: linalg.generic {{.*}} ins(%[[RAW]], +util.func public @propagate_unset_encoding_through_generic_with_scalar(%arg0: tensor<4096x?xf32, #encoding>, %arg1: f32, %arg2: index) -> tensor<4096x?xf32> { + %1 = iree_encoding.unset_encoding %arg0 : tensor<4096x?xf32, #encoding> -> tensor<4096x?xf32>{%arg2} + %2 = tensor.empty(%arg2) : tensor<4096x?xf32> + %3 = linalg.generic {indexing_maps = [#map3, #map4, #map3], iterator_types = ["parallel", "parallel"]} ins(%1, %arg1 : tensor<4096x?xf32>, f32) outs(%2 : tensor<4096x?xf32>) { + ^bb0(%in: f32, %in_0: f32, %out: f32): + %4 = arith.mulf %in, %in_0 : f32 + linalg.yield %4 : f32 + } -> tensor<4096x?xf32> + util.return %3 : tensor<4096x?xf32> +} + +// ----- + +#map_sink = affine_map<(d0, d1, d2) -> (d0, d2)> +#map_sink1 = affine_map<(d0, d1, d2) -> (d1, d2)> +#map_sink2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#map_sink3 = affine_map<(d0, d1) -> (d0, d1)> +#map_sink4 = affine_map<(d0, d1) -> ()> +#encoding_sink = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#map_sink, #map_sink1, #map_sink2], iteration_sizes = [?, 4096, 4096]> +// CHECK-LABEL: @sink_unset_encoding_with_encoding_dims( +// CHECK: %[[SCALE:.*]] = iree_encoding.set_encoding %arg1 encoding_dims{%arg2} +// CHECK: %[[RESULT:.*]] = linalg.generic {{.*}} ins(%arg0, %[[SCALE]] +// CHECK-SAME: tensor<?x4096xbf16, # +// CHECK: iree_encoding.unset_encoding %[[RESULT]] encoding_dims{%arg2} +util.func public @sink_unset_encoding_with_encoding_dims(%arg0: tensor<?x4096xf32, #encoding_sink>, %arg1: tensor<f32>, %m: index) -> tensor<?x4096xbf16> { + %1 = iree_encoding.unset_encoding %arg0 encoding_dims{%m} : tensor<?x4096xf32, #encoding_sink> -> tensor<?x4096xf32>{%m} + %2 = tensor.empty(%m) : tensor<?x4096xbf16> + %3 = linalg.generic {indexing_maps = [#map_sink3, #map_sink4, #map_sink3], iterator_types = ["parallel", "parallel"]} ins(%1, %arg1 : tensor<?x4096xf32>, tensor<f32>) outs(%2 : tensor<?x4096xbf16>) { + ^bb0(%in: f32, %in_0: f32, %out: bf16): + %4 = arith.mulf %in, %in_0 : f32 + %5 = arith.truncf %4 : f32 to bf16 + linalg.yield %5 : bf16 + } -> tensor<?x4096xbf16> + util.return %3 : tensor<?x4096xbf16> +} + +// ----- + +// Encoding ops whose producers are in the same dispatch region are left to +// hoist-encoding-ops. +#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d1, d2)> +#encoding = #iree_encoding.encoding<operand_index = 1 : 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, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]> +// CHECK-LABEL: @no_bubble_inside_existing_dispatch( +// CHECK-NOT: iree_encoding.set_encoding +// CHECK: flow.dispatch.region +// CHECK: %[[BCAST:.*]] = linalg.generic +// CHECK: iree_encoding.set_encoding %[[BCAST]] +// CHECK: flow.return +util.func public @no_bubble_inside_existing_dispatch( + %arg0: tensor<11008x128xf32>) -> tensor<2x11008x128xf32, #encoding> { + %6 = flow.dispatch.region -> (tensor<2x11008x128xf32, #encoding>) { + %8 = tensor.empty() : tensor<2x11008x128xf32> + %11 = linalg.generic + {indexing_maps = [#map1, #map], + iterator_types = ["parallel", "parallel", "parallel"]} + ins(%arg0 : tensor<11008x128xf32>) + outs(%8 : tensor<2x11008x128xf32>) { + ^bb0(%in: f32, %out: f32): + linalg.yield %in : f32 + } -> tensor<2x11008x128xf32> + %13 = iree_encoding.set_encoding %11 : tensor<2x11008x128xf32> -> tensor<2x11008x128xf32, #encoding> + flow.return %13 : tensor<2x11008x128xf32, #encoding> + } + util.return %6 : tensor<2x11008x128xf32, #encoding> +} + +// ----- + +// The broadcast producer is outside the dispatch region that holds the +// set_encoding, so the set_encoding is not bubbled out of the dispatch. +#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d1, d2)> +#encoding = #iree_encoding.encoding<operand_index = 1 : 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, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]> +// CHECK-LABEL: @no_bubble_out_of_dispatch( +// CHECK-NOT: iree_encoding.set_encoding +// CHECK: %[[BCAST:.+]] = linalg.generic +// CHECK: flow.dispatch.region +// CHECK-NEXT: iree_encoding.set_encoding %[[BCAST]] +util.func public @no_bubble_out_of_dispatch(%arg0: tensor<11008x128xf32>) -> tensor<2x11008x128xf32, #encoding> { + %empty = tensor.empty() : tensor<2x11008x128xf32> + %bcast = linalg.generic { + indexing_maps = [#map1, #map], + iterator_types = ["parallel", "parallel", "parallel"]} + ins(%arg0 : tensor<11008x128xf32>) + outs(%empty : tensor<2x11008x128xf32>) { + ^bb0(%in: f32, %out: f32): + linalg.yield %in : f32 + } -> tensor<2x11008x128xf32> + %result = flow.dispatch.region -> (tensor<2x11008x128xf32, #encoding>) { + %encoded = iree_encoding.set_encoding %bcast : tensor<2x11008x128xf32> -> tensor<2x11008x128xf32, #encoding> + %init = tensor.empty() : tensor<2x11008x128xf32, #encoding> + %neg = linalg.generic { + indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel", "parallel"]} + ins(%encoded : tensor<2x11008x128xf32, #encoding>) + outs(%init : tensor<2x11008x128xf32, #encoding>) { + ^bb0(%in: f32, %out: f32): + %0 = arith.negf %in : f32 + linalg.yield %0 : f32 + } -> tensor<2x11008x128xf32, #encoding> + flow.return %neg : tensor<2x11008x128xf32, #encoding> + } + util.return %result : tensor<2x11008x128xf32, #encoding> +} + +// ----- + +// The only consumer is inside a dispatch region, so the unset_encoding outside +// of it is not sunk into the dispatch. +#map = affine_map<(d0, d1, d2) -> (d0, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d1, d2)> +#map2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#map3 = affine_map<(d0, d1) -> (d0, d1)> +#map4 = affine_map<(d0, d1) -> ()> +#encoding = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2]> +// CHECK-LABEL: @no_sink_into_dispatch( +// CHECK: %[[RAW:.+]] = iree_encoding.unset_encoding +// CHECK: flow.dispatch.region +// CHECK: linalg.generic {{.*}} ins(%[[RAW]], %{{.+}} : tensor<4096x?xf32>, f32) +// CHECK-NOT: iree_encoding.unset_encoding +// CHECK: util.return +util.func public @no_sink_into_dispatch(%arg0: tensor<4096x?xf32, #encoding>, %arg1: f32, %arg2: index) -> tensor<4096x?xf32> { + %raw = iree_encoding.unset_encoding %arg0 : tensor<4096x?xf32, #encoding> -> tensor<4096x?xf32>{%arg2} + %result = flow.dispatch.region -> (tensor<4096x?xf32>{%arg2}) { + %empty = tensor.empty(%arg2) : tensor<4096x?xf32> + %scaled = linalg.generic {indexing_maps = [#map3, #map4, #map3], iterator_types = ["parallel", "parallel"]} ins(%raw, %arg1 : tensor<4096x?xf32>, f32) outs(%empty : tensor<4096x?xf32>) { + ^bb0(%in: f32, %in_0: f32, %out: f32): + %mul = arith.mulf %in, %in_0 : f32 + linalg.yield %mul : f32 + } -> tensor<4096x?xf32> + flow.return %scaled : tensor<4096x?xf32> + } + util.return %result : tensor<4096x?xf32> +} + +// ----- + +// The broadcast producer is outside the scf.if that holds the set_encoding. +// Propagation stops at the region boundary even though the rewrite would be +// valid inside the region. +#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d1, d2)> +#encoding = #iree_encoding.encoding<operand_index = 1 : 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, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]> +// CHECK-LABEL: @no_bubble_across_region( +// CHECK-NOT: iree_encoding.set_encoding +// CHECK: %[[BCAST:.+]] = linalg.generic +// CHECK: scf.if +// CHECK-NEXT: iree_encoding.set_encoding %[[BCAST]] +util.func public @no_bubble_across_region(%arg0: tensor<11008x128xf32>, %init: tensor<2x11008x128xf32, #encoding>, %cond: i1) -> tensor<2x11008x128xf32, #encoding> { + %empty = tensor.empty() : tensor<2x11008x128xf32> + %bcast = linalg.generic { + indexing_maps = [#map1, #map], + iterator_types = ["parallel", "parallel", "parallel"]} + ins(%arg0 : tensor<11008x128xf32>) + outs(%empty : tensor<2x11008x128xf32>) { + ^bb0(%in: f32, %out: f32): + linalg.yield %in : f32 + } -> tensor<2x11008x128xf32> + %result = scf.if %cond -> (tensor<2x11008x128xf32, #encoding>) { + %encoded = iree_encoding.set_encoding %bcast : tensor<2x11008x128xf32> -> tensor<2x11008x128xf32, #encoding> + scf.yield %encoded : tensor<2x11008x128xf32, #encoding> + } else { + scf.yield %init : tensor<2x11008x128xf32, #encoding> + } + util.return %result : tensor<2x11008x128xf32, #encoding> +} + +// ----- + +// The only consumer is inside an scf.if, so the unset_encoding outside of it is +// not sunk into the region, even though the rewrite would be valid. +#map = affine_map<(d0, d1, d2) -> (d0, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d1, d2)> +#map2 = affine_map<(d0, d1, d2) -> (d0, d1)> +#map3 = affine_map<(d0, d1) -> (d0, d1)> +#map4 = affine_map<(d0, d1) -> ()> +#encoding = #iree_encoding.encoding<operand_index = 2 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2]> +// CHECK-LABEL: @no_sink_across_region( +// CHECK: %[[RAW:.+]] = iree_encoding.unset_encoding +// CHECK: scf.if +// CHECK: linalg.generic {{.*}} ins(%[[RAW]], %{{.+}} : tensor<4096x?xf32>, f32) +// CHECK-NOT: iree_encoding.unset_encoding +// CHECK: util.return +util.func public @no_sink_across_region(%arg0: tensor<4096x?xf32, #encoding>, %arg1: f32, %arg2: index, %cond: i1) -> tensor<4096x?xf32> { + %raw = iree_encoding.unset_encoding %arg0 : tensor<4096x?xf32, #encoding> -> tensor<4096x?xf32>{%arg2} + %empty = tensor.empty(%arg2) : tensor<4096x?xf32> + %result = scf.if %cond -> (tensor<4096x?xf32>) { + %scaled = linalg.generic {indexing_maps = [#map3, #map4, #map3], iterator_types = ["parallel", "parallel"]} ins(%raw, %arg1 : tensor<4096x?xf32>, f32) outs(%empty : tensor<4096x?xf32>) { + ^bb0(%in: f32, %in_0: f32, %out: f32): + %mul = arith.mulf %in, %in_0 : f32 + linalg.yield %mul : f32 + } -> tensor<4096x?xf32> + scf.yield %scaled : tensor<4096x?xf32> + } else { + scf.yield %empty : tensor<4096x?xf32> + } + util.return %result : tensor<4096x?xf32> +}