[mlir] fix out-of-bounds in reduction tiling

[ftynse]

Summary

A transformation tiling a reduction dimension of a Linalg op needs a
tile size for said dimension. When an insufficient number of dimensions
was provided, it would segfault due to out-of-bounds access to a vector.

Also fix incorrect error reporting in the structured transform op
exercising this functionality.

[springerm]

mlir/test/Dialect/Linalg/transform-tile-reduction.mlir

309

typo?

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

transform::TileReductionUsingForeachThreadOp::applyToOne(
   SmallVector<OpFoldResult> tileSizes =
       getAsOpFoldResult(rewriter.getI64ArrayAttr(getTileSizes()));
   FailureOr<linalg::ForeachThreadReductionTilingResult> result =
       linalg::tileReductionUsingForeachThread(
           rewriter, cast<PartialReductionOpInterface>(target.getOperation()),
           numThreads, tileSizes, getMapping());
 
   if (failed(result)) {
-    results.assign(3, nullptr);
-    Diagnostic diag(target->getLoc(), DiagnosticSeverity::Remark);
-    diag << "could not tile reduction in target.";
-    return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
+    results.assign(4, nullptr);
+    auto diag = emitSilenceableError() << "could not tile reduction";
+    diag.attachNote(target.getLoc()) << "target operation";
+    return diag;
   }
   results.push_back(result->loops);
   results.push_back(result->initialOp);
   results.push_back(result->parallelTiledOp);
   results.push_back(result->mergeOp);
   return DiagnosedSilenceableFailure::success();
 }
 

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

linalg::tileReductionUsingForeachThread(RewriterBase &b,
   if (redDims.size() != 1)
     return b.notifyMatchFailure(
         op, "only support ops with one reduction dimension.");
   if (!tileSizes.empty() && tileSizes.size() != numThreads.size())
     return b.notifyMatchFailure(op, "if tile sizes are present it must have as "
                                     "many elements as number of threads");
   int reductionDim = static_cast<int>(redDims.front());
 
+  if (redDims.front() >= numThreads.size())
+    return b.notifyMatchFailure(
+        op, "reduction dimension must be mapped to threads");
+
   // 1. Create the inital tensor value.
   FailureOr<Operation *> identityTensor =
       op.generateInitialTensorForPartialReduction(b, loc, numThreads,
                                                   reductionDim);
   if (failed(identityTensor))
     return b.notifyMatchFailure(op,
                                 "cannot create a tensor of identity value.");
 

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

FailureOr<Operation *> generateInitialTensorForPartialReduction(
     OpBuilder::InsertionGuard guard(b);
     assert(reductionDims.size() == 1 &&
            "only support single reduction right now.");
     if (linalgOp.hasBufferSemantics())
       return op->emitOpError("expected operation to have tensor semantics");
     // Insert the new parallel dimension based on the index of the reduction
     // loop. This could be controlled by user for more flexibility.
     int64_t insertSplitDimension = reductionDims[0];
+    assert(sizes.size() >= static_cast<size_t>(insertSplitDimension) &&
+           "reduction dimension must be tiled");
 
     SmallVector<Operation *, 4> combinerOps;
     if (!matchReduction(linalgOp.getRegionOutputArgs(), 0, combinerOps) ||
         combinerOps.size() != 1)
       return op->emitOpError("Failed to anaysis the reduction operation.");
 
     Operation *reductionOp = combinerOps[0];
     Optional<Attribute> identity = getNeutralElement(reductionOp);

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

mlir::scf::tileReductionUsingScf(PatternRewriter &b,
   int reductionDim;
   for (auto &[idx, iteratorType] :
        llvm::enumerate(tilingInterfaceOp.getLoopIteratorTypes())) {
     if (iteratorType == utils::IteratorType::reduction) {
       reductionDim = idx;
       break;
     }
   }
+  if (static_cast<size_t>(reductionDim) >= tileSize.size())
+    return b.notifyMatchFailure(op, "reduction dimension must be tiled");
+
   // 1. create the inital tensor value.
   FailureOr<Operation *> identityTensor =
       op.generateInitialTensorForPartialReduction(b, loc, tileSize,
                                                   reductionDim);
   if (failed(identityTensor))
     return b.notifyMatchFailure(op,
                                 "cannot create a tensor of identity value.");
   // 2. Create the nested loops.

mlir/test/Dialect/Linalg/transform-tile-reduction.mlir

-// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -canonicalize -cse | FileCheck %s
+// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -canonicalize -cse -verify-diagnostics | FileCheck %s
 
 func.func @reduction_tile(%arg0: tensor<?x?xf32>, %out: tensor<?xf32>) -> tensor<?xf32> {
   %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
                                           affine_map<(d0, d1) -> (d0)>],
    iterator_types = ["parallel", "reduction"]}
    ins(%arg0 : tensor<?x?xf32>)
    outs(%out : tensor<?xf32>) {
     ^bb0(%arg7: f32, %arg9: f32):
^bb0(%arg1: !pdl.operation):
   // CHECK-NEXT:     linalg.generic
   //      CHECK:     iterator_types = ["parallel", "reduction"]
   transform.print %2 {name = "expecting parallel reduction"} : !pdl.operation
   //      CHECK:     expecting parallel reduction
   // CHECK-NEXT:     linalg.generic
   //      CHECK:     iterator_types = ["parallel", "reduction"]
   transform.print %3 {name = "expecting parallel reduction"} : !pdl.operation
 }
+
+// -----
+
+func.func @reduction_untiled_foreach_thread(
+  %arg0: tensor<?x?xf32>, %out: tensor<?xf32>) -> tensor<?xf32> {
+  // expected-note @below {{target operation}}
+  %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                                          affine_map<(d0, d1) -> (d0)>],
+   iterator_types = ["parallel", "reduction"]}
+   ins(%arg0 : tensor<?x?xf32>)
+   outs(%out : tensor<?xf32>) {
+    ^bb0(%arg7: f32, %arg9: f32):
+      %1 = arith.mulf %arg7, %arg7 : f32
+      %2 = arith.addf %1, %arg9 : f32
+      linalg.yield %2 : f32
+    } -> tensor<?xf32>
+  return %red : tensor<?xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
+  // expected-error @below {{could not tile reduction}}
+  %loop, %1, %2, %3 = transform.structured.tile_reduction_using_foreach_thread %0
+    by num_threads = [5], tile_sizes = [3], mapping = [#gpu.thread<x>]
+
+}