Populate outer parallel loops for UnsetEncodingOp (#12289)

Similar to `SetEncodingOp`, this change populates `the outer parallel
loops = result rank` for `UnsetEncodingOp`.

Fix #12285

compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp
index aa0a30b..90b671a 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.cpp
@@ -251,6 +251,10 @@
   if (auto setEncodingOp = dyn_cast<IREE::LinalgExt::SetEncodingOp>(op)) {
     return llvm::SmallBitVector(setEncodingOp.getResultType().getRank(), true);
   }
+  if (auto unsetEncodingOp = dyn_cast<IREE::LinalgExt::UnsetEncodingOp>(op)) {
+    return llvm::SmallBitVector(unsetEncodingOp.getResultType().getRank(),
+                                true);
+  }
 
   auto interfaceOp = dyn_cast<TilingInterface>(op);
   if (!interfaceOp) {

compiler/src/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir
index b9c6293..2f69d54 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/dispatch_linalg_on_tensors.mlir
@@ -2004,6 +2004,55 @@
 
 // -----
 
+#map = affine_map<(d0, d1) -> (d1)>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+module {
+  func.func @root_on_unset_encoding(
+      %arg0: tensor<784x96xf32, #iree_linalg_ext.encoding<user = MATMUL_F32F32F32, role = LHS>>,
+      %arg1: tensor<96xf32>) -> tensor<784x96xf32> {
+    %0 = iree_linalg_ext.unset_encoding %arg0 : tensor<784x96xf32, #iree_linalg_ext.encoding<user = MATMUL_F32F32F32, role = LHS>> -> tensor<784x96xf32>
+    %1 = tensor.empty() : tensor<784x96xf32>
+    %cst = arith.constant 0.000000e+00 : f32
+    %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<784x96xf32>) -> tensor<784x96xf32>
+    %3 = linalg.generic {indexing_maps = [#map, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%arg1, %0 : tensor<96xf32>, tensor<784x96xf32>) outs(%1 : tensor<784x96xf32>) {
+    ^bb0(%in: f32, %in_0: f32, %out: f32):
+      %5 = arith.addf %in, %in_0 : f32
+      linalg.yield %5 : f32
+    } -> tensor<784x96xf32>
+    %4 = linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%3, %3 : tensor<784x96xf32>, tensor<784x96xf32>) outs(%1 : tensor<784x96xf32>) {
+    ^bb0(%in: f32, %in_0: f32, %out: f32):
+      %5 = arith.addf %in, %in_0 : f32
+      linalg.yield %5 : f32
+    } -> tensor<784x96xf32>
+    return %4 : tensor<784x96xf32>
+  }
+}
+//      CHECL: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d1)>
+//      CHECK: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+//      CHECK: func @root_on_unset_encoding
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<784x96xf32, #iree_linalg_ext.encoding<user = MATMUL_F32F32F32, role = LHS>>
+// CHECK-SAME:     %[[ARG1:.+]]: tensor<96xf32>
+//      CHECK:   %[[DISPATCH:.+]] = flow.dispatch.workgroups(%[[ARG0]], %[[ARG1]])
+// CHECK-NEXT:     %[[INARG0:.+]]: !flow.dispatch.tensor<readonly:tensor<784x96xf32, #iree_linalg_ext.encoding<user = MATMUL_F32F32F32, role = LHS>>>
+// CHECK-SAME:     %[[INARG1:.+]]: !flow.dispatch.tensor<readonly:tensor<96xf32>>
+// CHECK-SAME:     %[[OUTARG:.+]]: !flow.dispatch.tensor<writeonly:tensor<784x96xf32>>
+//      CHECK:     %[[LOAD0:.+]] = flow.dispatch.tensor.load %[[INARG0]], offsets = [0, 0], sizes = [784, 96], strides = [1, 1]
+// CHECK-SAME:       : !flow.dispatch.tensor<readonly:tensor<784x96xf32, #iree_linalg_ext.encoding<user = MATMUL_F32F32F32, role = LHS>>>
+//      CHECK:     %[[LOAD1:.+]] = flow.dispatch.tensor.load %[[INARG1]], offsets = [0], sizes = [96], strides = [1]
+// CHECK-SAME:       : !flow.dispatch.tensor<readonly:tensor<96xf32>>
+//      CHECK:     %[[OUT:.+]] = tensor.empty() : tensor<784x96xf32>
+//      CHECK:     %[[UNSET:.+]] = iree_linalg_ext.unset_encoding %[[LOAD0]] : tensor<784x96xf32, #iree_linalg_ext.encoding<user = MATMUL_F32F32F32, role = LHS>> -> tensor<784x96xf32>
+//      CHECK:     %[[GENERIC0:.+]] = linalg.generic {indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP1]]],
+// CHECK-SAME:       iterator_types = ["parallel", "parallel"]} ins(%[[LOAD1]], %[[UNSET]]
+// CHECK-SAME:       : tensor<96xf32>, tensor<784x96xf32>) outs(%[[OUT]] : tensor<784x96xf32>)
+//      CHECK:     %[[GENERIC1:.+]] = linalg.generic {indexing_maps = [#[[MAP1]], #[[MAP1]], #[[MAP1]]],
+// CHECK-SAME:       iterator_types = ["parallel", "parallel"]} ins(%[[GENERIC0]], %[[GENERIC0]]
+// CHECK-SAME:       : tensor<784x96xf32>, tensor<784x96xf32>) outs(%[[OUT]] : tensor<784x96xf32>)
+//      CHECK:     flow.dispatch.tensor.store %[[GENERIC1]], %[[OUTARG]], offsets = [0, 0], sizes = [784, 96], strides = [1, 1] : tensor<784x96xf32>
+//      CHECK:     flow.return
+
+// -----
+
 func.func @gemm_encoded(
     %arg0 : tensor<?x?xf32, #iree_linalg_ext.encoding<user = MATMUL_F32F32F32, role = LHS>>,
     %arg1 : tensor<?x?xf32, #iree_linalg_ext.encoding<user = MATMUL_F32F32F32, role = RHS>>,