[mlir][linalg] Expand test case for tile-and-fuse with transform dialect

[springerm]

Summary

Reverse the order of the payload ops. fuse_into_containing_op should still work.

[ftynse]

We may eventually want a fuzzer that randomizes the order.

mlir/test/Dialect/Linalg/transform-tile-and-fuse.mlir

-// RUN: mlir-opt %s --test-transform-dialect-interpreter -canonicalize | FileCheck %s
+// RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file -canonicalize | FileCheck %s
 
 // This is a simple tile-and-fuse example with a single fusion group.
 
 module {
   // CHECK: func @foo
   // CHECK:   scf.foreach_thread {{.*}} {
   // CHECK:     linalg.fill
   // CHECK:     linalg.matmul
   // CHECK:     linalg.generic
   // CHECK:   }
   func.func @foo(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?xf32>,
                  %D: tensor<?x?xf32>, %sz0: index, %sz1: index)
       -> tensor<?x?xf32>
   {
     %cst = arith.constant 0.000000e+00 : f32
     %5 = linalg.fill
         {__producer__}
         ins(%cst : f32)
         outs(%D : tensor<?x?xf32>) -> tensor<?x?xf32>
     %6 = linalg.matmul
         {__producer__}
         ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
         outs(%5 : tensor<?x?xf32>) -> tensor<?x?xf32>
-    %7 = linalg.generic 
+    %7 = linalg.generic
         {__root__,
          indexing_maps = [affine_map<(d0, d1) -> (d0)>,
                           affine_map<(d0, d1) -> (d0, d1)>,
                           affine_map<(d0, d1) -> (d0, d1)>],
          iterator_types = ["parallel", "parallel"]
         }
         ins(%C, %6 : tensor<?xf32>, tensor<?x?xf32>)
         outs(%D : tensor<?x?xf32>) {
^bb1(%arg1: !pdl.operation):
       // Tile the root.
       %foreach_thread_op, %tiled_op = transform.structured.tile_to_foreach_thread_op %root num_threads [10, 20]
 
       // Fuse all producers.
       transform.structured.fuse_into_containing_op %producers into %foreach_thread_op
     }
   }
 }
+
+// -----
+
+// Inverse the order of the payload ops passed to the tile_to_foreach_thread_op
+// op. Fusion should still work.
+
+module {
+  // CHECK: func @foo
+  // CHECK:   scf.foreach_thread {{.*}} {
+  // CHECK:     linalg.fill
+  // CHECK:     linalg.matmul
+  // CHECK:     linalg.generic
+  // CHECK:   }
+  func.func @foo(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?xf32>,
+                 %D: tensor<?x?xf32>, %sz0: index, %sz1: index)
+      -> tensor<?x?xf32>
+  {
+    %cst = arith.constant 0.000000e+00 : f32
+    %5 = linalg.fill
+        {__producer__}
+        ins(%cst : f32)
+        outs(%D : tensor<?x?xf32>) -> tensor<?x?xf32>
+    %6 = linalg.matmul
+        {__producer__}
+        ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
+        outs(%5 : tensor<?x?xf32>) -> tensor<?x?xf32>
+    %7 = linalg.generic
+        {__root__,
+         indexing_maps = [affine_map<(d0, d1) -> (d0)>,
+                          affine_map<(d0, d1) -> (d0, d1)>,
+                          affine_map<(d0, d1) -> (d0, d1)>],
+         iterator_types = ["parallel", "parallel"]
+        }
+        ins(%C, %6 : tensor<?xf32>, tensor<?x?xf32>)
+        outs(%D : tensor<?x?xf32>) {
+    ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
+      %16 = arith.maxf %arg3, %cst : f32
+      %17 = arith.cmpf ogt, %arg2, %cst : f32
+      %18 = arith.select %17, %cst, %16 : f32
+      linalg.yield %18 : f32
+    } -> tensor<?x?xf32>
+    return %7 : tensor<?x?xf32>
+  }
+
+  transform.with_pdl_patterns {
+  ^bb0(%arg0: !pdl.operation):
+    transform.sequence %arg0 {
+    ^bb1(%arg1: !pdl.operation):
+      // Find the root and all producers.
+      %root = transform.structured.match attribute{"__root__"} in %arg1
+      %producers = transform.structured.match attribute{"__producer__"} in %arg1
+      %reversed_producers = transform.test_reverse_payload_ops %producers
+
+      // Tile the root.
+      %foreach_thread_op, %tiled_op = transform.structured.tile_to_foreach_thread_op %root num_threads [10, 20]
+
+      // Fuse all producers.
+      transform.structured.fuse_into_containing_op %reversed_producers into %foreach_thread_op
+    }
+  }
+}

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

DiagnosedSilenceableFailure mlir::test::TestRemapOperandPayloadToSelfOp::apply(
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure mlir::test::TestRemoveTestExtensionOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   state.removeExtension<TestTransformStateExtension>();
   return DiagnosedSilenceableFailure::success();
 }
+
+DiagnosedSilenceableFailure
+mlir::test::TestReversePayloadOpsOp::apply(transform::TransformResults &results,
+                                           transform::TransformState &state) {
+  ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
+  auto reversedOps = llvm::to_vector(llvm::reverse(payloadOps));
+  results.set(getResult().cast<OpResult>(), reversedOps);
+  return DiagnosedSilenceableFailure::success();
+}
+
 DiagnosedSilenceableFailure mlir::test::TestTransformOpWithRegions::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   return DiagnosedSilenceableFailure::success();
 }
 
 void mlir::test::TestTransformOpWithRegions::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {}
 

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.td

 def TestRemoveTestExtensionOp
   : Op<Transform_Dialect, "test_remove_test_extension",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         NoSideEffect]> {
   let assemblyFormat = "attr-dict";
   let cppNamespace = "::mlir::test";
 }
 
+def TestReversePayloadOpsOp
+  : Op<Transform_Dialect, "test_reverse_payload_ops",
+    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+     DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$result);
+  let assemblyFormat = "$target attr-dict";
+  let cppNamespace = "::mlir::test";
+}
+
 def TestTransformOpWithRegions
   : Op<Transform_Dialect, "test_transform_op_with_regions",
        [DeclareOpInterfaceMethods<TransformOpInterface>,
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let regions = (region AnyRegion:$first, AnyRegion:$second);
   let assemblyFormat = "attr-dict-with-keyword regions";
   let cppNamespace = "::mlir::test";
 }