address review comments

Author: tkarna

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

@@ -400,7 +400,9 @@ def FuseOp : Op<Transform_Dialect, "structured.fuse",
         TransformOpInterface, ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     Tiles the operations pointed to by the target handle and fuses their
-    producers greedily using the options provided as attributes.
+    producers greedily using the options provided as attributes. Tile sizes
+    and loop interchange permutation can be provided as either static
+    attributes or dynamic values (transform parameters or payload handles).
 
     If `apply_cleanup` is true then slice canonicalization is applied between
     fusion steps. If `use_forall` is true then tiling method generates a

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

@@ -759,14 +759,21 @@ transform::FuseOp::apply(transform::TransformRewriter &rewriter,
 }
 
 LogicalResult transform::FuseOp::verify() {
+  auto iterspace_dim = getStaticTileSizes().size();
   ArrayRef<int64_t> permutation = getStaticTileInterchange();
-  if (!llvm::any_of(permutation,
-                    [](int64_t v) { return ShapedType::isDynamic(v); })) {
-    auto sequence = llvm::to_vector(llvm::seq<int64_t>(0, permutation.size()));
-    if (!std::is_permutation(sequence.begin(), sequence.end(),
-                             permutation.begin(), permutation.end())) {
-      return emitOpError() << "expects interchange to be a permutation, found "
-                           << getTileInterchange();
+  if (permutation.size() > iterspace_dim)
+    return emitOpError()
+           << "interchange length exceeds iteration space dimensions ("
+           << iterspace_dim << "), found " << getTileInterchange();
+  llvm::SmallDenseSet<int64_t, 4> seen;
+  for (int64_t v : permutation) {
+    if (!ShapedType::isDynamic(v)) {
+      if (v < 0 || v >= iterspace_dim)
+        return emitOpError() << "expects interchange values to be in range [0, "
+                             << iterspace_dim << "), found: " << v;
+      auto result = seen.insert(v);
+      if (!result.second)
+        return emitOpError() << "found duplicate interchange value: " << v;
     }
   }
 
@@ -780,37 +787,12 @@ LogicalResult transform::FuseOp::verify() {
 }
 
 SmallVector<OpFoldResult> transform::FuseOp::getMixedTileSizes() {
-  ValueRange dynamicValues = getTileSizes();
-  ArrayRef<int64_t> staticValues = getStaticTileSizes();
-  SmallVector<OpFoldResult> results;
-  results.reserve(staticValues.size());
-  unsigned dynamicPos = 0;
-  Builder builder(getContext());
-  for (int64_t size : staticValues) {
-    if (size == ShapedType::kDynamic) {
-      results.push_back(dynamicValues[dynamicPos++]);
-    } else {
-      results.push_back(builder.getIndexAttr(size));
-    }
-  }
-  return results;
+  return getMixedValues(getStaticTileSizes(), getTileSizes(), getContext());
 }
 
 SmallVector<OpFoldResult> transform::FuseOp::getMixedTileInterchange() {
-  ValueRange dynamicValues = getTileInterchange();
-  ArrayRef<int64_t> staticValues = getStaticTileInterchange();
-  SmallVector<OpFoldResult> results;
-  results.reserve(staticValues.size());
-  unsigned dynamicPos = 0;
-  Builder builder(getContext());
-  for (int64_t size : staticValues) {
-    if (size == ShapedType::kDynamic) {
-      results.push_back(dynamicValues[dynamicPos++]);
-    } else {
-      results.push_back(builder.getIndexAttr(size));
-    }
-  }
-  return results;
+  return getMixedValues(getStaticTileInterchange(), getTileInterchange(),
+                        getContext());
 }
 
 void transform::FuseOp::getEffects(

mlir/python/mlir/dialects/transform/structured.py

@@ -146,8 +146,8 @@ def __init__(
         *,
         tile_sizes: Optional[MixedValues] = None,
         tile_interchange: Optional[MixedValues] = None,
-        apply_cleanup: Optional[bool] = False,
-        use_forall: Optional[bool] = False,
+        apply_cleanup: bool = False,
+        use_forall: bool = False,
         loc=None,
         ip=None,
     ):
@@ -160,8 +160,8 @@ def __init__(
         *,
         tile_sizes: Optional[MixedValues] = None,
         tile_interchange: Optional[MixedValues] = None,
-        apply_cleanup: Optional[bool] = False,
-        use_forall: Optional[bool] = False,
+        apply_cleanup: bool = False,
+        use_forall: bool = False,
         loc=None,
         ip=None,
     ):
@@ -174,8 +174,8 @@ def __init__(
         *,
         tile_sizes: Optional[MixedValues] = None,
         tile_interchange: Optional[MixedValues] = None,
-        apply_cleanup: Optional[bool] = False,
-        use_forall: Optional[bool] = False,
+        apply_cleanup: bool = False,
+        use_forall: bool = False,
         loc=None,
         ip=None,
     ):
@@ -192,7 +192,7 @@ def __init__(
             _,
         ) = _dispatch_dynamic_index_list(tile_interchange)
         num_loops = (
-            1 if use_forall else sum(0 if v == 0 else 1 for v in static_tile_sizes)
+            1 if use_forall else sum(1 for v in static_tile_sizes if v != 0)
         )
 
         if isinstance(loop_types_or_target, (Operation, Value, OpView)):

mlir/test/python/dialects/transform_structured_ext.py

@@ -158,6 +158,26 @@ def testFuseOpParams(target):
     # CHECK-SAME: (!transform.any_op, !transform.param<i64>, !transform.param<i64>) -> (!transform.any_op, !transform.any_op, !transform.any_op)
 
 
+@run
+@create_sequence
+def testFuseOpHandles(target):
+    size1 = structured.MatchOp.match_op_names(target, ["arith.constant"])
+    ichange1 = structured.MatchOp.match_op_names(target, ["arith.constant"])
+    structured.FuseOp(
+        target,
+        tile_sizes=[size1, 8],
+        tile_interchange=[ichange1, 1],
+    )
+    # CHECK-LABEL: TEST: testFuseOpHandles
+    # CHECK: transform.sequence
+    # CHECK: %[[H:.*]] = transform.structured.match
+    # CHECK: %[[I:.*]] = transform.structured.match
+    # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.fuse
+    # CHECK-SAME: tile_sizes [%[[H]], 8]
+    # CHECK-SAME: interchange [%[[I]], 1]
+    # CHECK-SAME: (!transform.any_op, !transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+
+
 @run
 @create_sequence
 def testFuseOpAttributes(target):