[mlir] improve error handling in Linalg op splitting

[ftynse]

Summary

In several cases, the splitting may be known to be a noop, i.e., produce
no second part. Thread this information through the transform utilities
to the transform dialect, and differentiate it from the error state.

[nicolasvasilache]

Thanks for fixing this!

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 /// Fuse two `linalg.generic` operations that have a producer-consumer
 /// relationship captured through `fusedOperand`. The method expects
 /// that `areElementwiseOpsFusable` returns true for the given `fusedOperand`.
 FailureOr<Operation *> fuseElementwiseOps(RewriterBase &rewriter,
                                           OpOperand *fusedOperand);
 
 /// Split the given `op` into two parts along the given iteration space
 /// `dimension` at the specified `splitPoint`, and return the two parts.
+/// If the second part is statically known to be empty, do not create it
+/// and return nullptr instead. Error state is signalled by returning
+/// a pair of nullptrs.
 ///
 /// For example, the following op:
 ///
 ///   linalg.matmul ins(%0, %1 : tensor<128x32xf32>, tensor<32x64xf32>)
 ///                 outs(%2 : tensor<128x64xf32>)
 ///
 /// split along the first dimension at position 42 will result in:
 ///

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

if (getDynamicSplitPoint()) {
     }
   } else {
     splitPoints.resize(payload.size(),
                        rewriter.getIndexAttr(getStaticSplitPoint()));
   }
 
   // Split each target operation.
   SmallVector<Operation *> first, second;
+  Operation *noSecondPart = nullptr;
   for (const auto &pair : llvm::zip(payload, splitPoints)) {
     Operation *target = std::get<0>(pair);
     auto linalgOp = dyn_cast<LinalgOp>(target);
     if (!linalgOp) {
       auto diag = emitSilenceableError() << "only applies to structured ops";
       diag.attachNote(target->getLoc()) << "target op";
       results.set(getFirst().cast<OpResult>(), {});
       results.set(getSecond().cast<OpResult>(), {});
       return diag;
     }
 
     if (getDimension() >= linalgOp.getNumLoops()) {
       auto diag = emitSilenceableError() << "dimension " << getDimension()
                                          << " does not exist in target op";
       diag.attachNote(target->getLoc()) << "target op";
       results.set(getFirst().cast<OpResult>(), {});
       results.set(getSecond().cast<OpResult>(), {});
       return diag;
     }
 
     rewriter.setInsertionPoint(linalgOp);
     std::tie(first.emplace_back(), second.emplace_back()) = linalg::splitOp(
         rewriter, cast<TilingInterface>(linalgOp.getOperation()),
         getDimension(), std::get<1>(pair));
+
+    // Propagate errors.
+    if (!first.back() && !second.back()) {
+      auto diag = emitDefiniteFailure() << "internal failure in splitting";
+      diag.attachNote(target->getLoc()) << "target op";
+      return diag;
+    }
+
+    // Do not add null second parts.
+    if (!second.back()) {
+      noSecondPart = target;
+      second.pop_back();
+    }
+  }
+
+  if (second.size() != first.size() && !second.empty()) {
+    results.set(getFirst().cast<OpResult>(), {});
+    results.set(getSecond().cast<OpResult>(), {});
+    auto diag =
+        emitSilenceableError()
+        << "splitting does not produce the second part for a subset of targets";
+    diag.attachNote() << "expected splitting to produce the second part of all "
+                         "or none of the targets";
+    diag.attachNote(noSecondPart->getLoc())
+        << "first target with no second part";
+    return diag;
   }
 
   results.set(getFirst().cast<OpResult>(), first);
   results.set(getSecond().cast<OpResult>(), second);
   return DiagnosedSilenceableFailure::success();
 }
 
 void SplitOp::getEffects(

mlir/lib/Dialect/Linalg/Transforms/Split.cpp

linalg::splitOp(RewriterBase &rewriter, TilingInterface op, unsigned dimension,
   // Create the second part.
   OpFoldResult totalOffset = makeComposedFoldedAffineApply(
       rewriter, op.getLoc(), d0 + d1, {offsets[dimension], minSplitPoint});
   SmallVector<Value> secondResults;
   TilingInterface secondPart =
       createSplitPart(rewriter, op.getLoc(), op, offsets, sizes, firstResults,
                       dimension, remainingSize, totalOffset, secondResults);
 
+  // Propagate any errors in part creation.
+  if (!firstPart || !secondPart)
+    return {TilingInterface(), TilingInterface()};
+
   // Replace the original op with the results of the two newly created ops.
   rewriter.replaceOp(op, secondResults);
   return {firstPart, secondPart};
 }

mlir/test/Dialect/Linalg/transform-op-split.mlir

^bb0(%0: f32, %1: f32):
     %call_res = func.call @elem(%0, %i, %i) : (f32, index, index) -> f32
     linalg.yield %call_res : f32
   } -> tensor<100xf32>
 
   // CHECK: return %[[RES]]
   return %0 : tensor<100xf32>
 }
 
+// -----
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
+  %1:2 = transform.structured.split %0 after 42 { dimension = 0 }
+}
+
+func.func private @elem(%arg0: f32, %arg1: index, %arg2: index) -> f32
+
 // CHECK-LABEL: @one_d_static_overflow
 // CHECK-SAME:  %[[IN:.+]]: tensor<10xf32>, %[[OUT:.+]]: tensor<10xf32>
 func.func @one_d_static_overflow(%arg0: tensor<10xf32>, %arg1: tensor<10xf32>) -> tensor<10xf32> {
   // Folding is sufficiently powerful to detect the static overflow and avoid
   // the splitting altogether.
   // CHECK: %[[RES_SLICE_LOW:.+]] = linalg.generic
   // CHECK:   ins(%[[IN]]
   // CHECK:   outs(%[[OUT]]
func.func @one_d_static(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
   }
   ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) {
   ^bb0(%0: f32, %1: f32):
     linalg.yield %0 : f32
   } -> tensor<100xf32>
   return %0 : tensor<100xf32>
 }
 
+// -----
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
+  // expected-error @below {{splitting does not produce the second part for a subset of targets}}
+  // expected-note @below {{expected splitting to produce the second part of all or none of the targets}}
+  %1:2 = transform.structured.split %0 after 142 { dimension = 0 }
+}
+
+func.func private @elem(%arg0: f32, %arg1: index, %arg2: index) -> f32
+
+func.func @split_one_but_not_other(
+    %arg0: tensor<100xf32>, %arg1: tensor<100xf32>,
+    %arg2: tensor<200xf32>, %arg3: tensor<200xf32>)
+    -> (tensor<100xf32>, tensor<200xf32>) {
+  // expected-note @below {{first target with no second part}}
+  %0 = linalg.generic {
+    indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
+    iterator_types = ["parallel"]
+  }
+  ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) {
+  ^bb0(%arg4: f32, %arg5: f32):
+    %i = linalg.index 0 : index
+    %call_res = func.call @elem(%arg4, %i, %i) : (f32, index, index) -> f32
+    linalg.yield %call_res : f32
+  } -> tensor<100xf32>
+
+  %1 = linalg.generic {
+    indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
+    iterator_types = ["parallel"]
+  }
+  ins(%arg2: tensor<200xf32>) outs(%arg3: tensor<200xf32>) {
+  ^bb0(%arg4: f32, %arg5: f32):
+    %i = linalg.index 0 : index
+    %call_res = func.call @elem(%arg4, %i, %i) : (f32, index, index) -> f32
+    linalg.yield %call_res : f32
+  } -> tensor<200xf32>
+
+  return %0, %1 : tensor<100xf32>, tensor<200xf32>
+}
+