[DispatchCreation] Do not collapse dimensions if there are encoded operands (#20621)

Signed-off-by: hanhanW <[email protected]>

Author: hanhanW

compiler/src/iree/compiler/DispatchCreation/CollapseDimensions.cpp

@@ -177,6 +177,14 @@ static bool isEligibleForCollapse(Operation *op) {
     return false;
   }
 
+  auto hasEncoding = [](Type type) -> bool {
+    auto rankedTensorType = dyn_cast<RankedTensorType>(type);
+    return rankedTensorType && rankedTensorType.getEncoding();
+  };
+  if (llvm::any_of(op->getOperandTypes(), hasEncoding)) {
+    return false;
+  }
+
   // TODO(guray) Currently we can only collapse when result of all the
   // AffineMaps are dimensions. Possible to collapse cases like
   // affine_map<d0, d1+d2> with affine_map<d0, d1+d2>, however, this is not

compiler/src/iree/compiler/DispatchCreation/test/collapse_dimensions.mlir

@@ -213,6 +213,26 @@ util.func public @quantized_matmul(%arg0: tensor<4096x32x128xi8>, %arg1: tensor<
 
 // -----
 
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+#encoding = #iree_encoding.testing_encoding<>
+util.func public @do_not_collapse_ops_with_encoding(%arg0: tensor<2x320x128x128xf32, #encoding>) -> tensor<2x320x128x128xf32, #encoding> {
+  %0 = flow.dispatch.region -> (tensor<2x320x128x128xf32, #encoding>) {
+    %empty = tensor.empty() : tensor<2x320x128x128xf32, #encoding>
+    %cst = arith.constant 3.14 : f32
+    %elementwise = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x320x128x128xf32, #encoding>) outs(%empty : tensor<2x320x128x128xf32, #encoding>) {
+    ^bb0(%in : f32, %out : f32):
+      %22 = arith.mulf %cst, %in : f32
+      linalg.yield %22 : f32
+    } -> tensor<2x320x128x128xf32, #encoding>
+    flow.return %elementwise : tensor<2x320x128x128xf32, #encoding>
+  }
+  util.return %0 : tensor<2x320x128x128xf32, #encoding>
+}
+// CHECK-LABEL: @do_not_collapse_ops_with_encoding(
+//   CHECK-NOT:   tensor.collapse_shape
+
+// -----
+
 #map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 util.func public @elementwise_chain(%arg0: tensor<2x320x128x128xf32>) -> tensor<2x320x128x128xf32> {
   %0 = flow.dispatch.region -> (tensor<2x320x128x128xf32>) {

tests/e2e/encoding/BUILD.bazel

@@ -17,7 +17,7 @@ iree_check_single_backend_test_suite(
     name = "check_rocm_hip",
     srcs = ["encoding.mlir"],
     compiler_flags = [
-        "--iree-global-opt-experimental-rocm-data-tiling",
+        "--iree-global-opt-enable-early-materialization=false",
         "--iree-hip-encoding-layout-resolver=data-tiling",
     ],
     driver = "hip",

tests/e2e/encoding/CMakeLists.txt

@@ -20,7 +20,7 @@ iree_check_single_backend_test_suite(
   DRIVER
     "hip"
   COMPILER_FLAGS
-    "--iree-global-opt-experimental-rocm-data-tiling"
+    "--iree-global-opt-enable-early-materialization=false"
     "--iree-hip-encoding-layout-resolver=data-tiling"
 )