[Mlir-commits] [mlir] faac898 - [mlir] fix out-of-bounds in reduction tiling

[Alex Zinenko]

Author: Alex Zinenko
Date: 2023-01-05T15:20:26Z
New Revision: faac8989871197492b7454cd1b259951a26b2f7a

URL: https://github.com/llvm/llvm-project/commit/faac8989871197492b7454cd1b259951a26b2f7a
DIFF: https://github.com/llvm/llvm-project/commit/faac8989871197492b7454cd1b259951a26b2f7a.diff

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

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.

Reviewed By: springerm, ThomasRaoux

Differential Revision: https://reviews.llvm.org/D141046

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
    mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
    mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
    mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
    mlir/test/Dialect/Linalg/transform-tile-reduction.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 5660891d56529..6702e9a057748 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -1276,10 +1276,10 @@ transform::TileReductionUsingForeachThreadOp::applyToOne(
           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);

diff  --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
index 7ffc5864e5f51..0c0e0d6e33970 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
@@ -626,6 +626,10 @@ linalg::tileReductionUsingForeachThread(RewriterBase &b,
                                     "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,

diff  --git a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
index cf8f56fcce068..9fcee2b7e0aad 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
@@ -258,6 +258,8 @@ struct LinalgOpPartialReductionInterface
     // 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) ||

diff  --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index 9d39a1309a5fe..02323c584a84f 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -424,6 +424,9 @@ mlir::scf::tileReductionUsingScf(PatternRewriter &b,
       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,

diff  --git a/mlir/test/Dialect/Linalg/transform-tile-reduction.mlir b/mlir/test/Dialect/Linalg/transform-tile-reduction.mlir
index 370930661bd20..1a23bcf887e7b 100644
--- a/mlir/test/Dialect/Linalg/transform-tile-reduction.mlir
+++ b/mlir/test/Dialect/Linalg/transform-tile-reduction.mlir
@@ -1,4 +1,4 @@
-// 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)>,
@@ -303,3 +303,30 @@ transform.sequence failures(propagate) {
   //      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>]
+
+}