[mlir] Flag near misses in file splitting

[jpienaar]

Summary

Flags some potential cases where splitting isn't happening and so could result
in confusing results. Also update some test files where there were near misses
in splitting that seemed unintentional.

[jpienaar]

Expand to flag larger and smaller mismatches based on feedback on cases seen.

[mehdi_amini]

mlir/lib/Support/ToolUtilities.cpp

49

This patch made me look at our current behavior: we don't require a new line after the pattern!
I think we should?

I mean that it is happy to parse this:

func private @foo()
// -----func private @foo()

->

$ ./bin/mlir-opt  -split-input-file split.mlir 
module  {
  func private @foo()
}

module  {
  func private @foo()
}

So I would check exactly --\n here.

51

Can you make splitMarker a StringRef from the beginning?

65

I think this loop deserve a bit more doc

[jpienaar]

Check for newline explicitly instead.

[rriddle]

mlir/lib/Support/ToolUtilities.cpp

49

We actually semi rely on this when we generate "print-ir-*" output. We generate output such that each output block can be treated as a separate IR block.

[rriddle]

What is the cost of doing these checks on large files?

[mehdi_amini]

mlir/lib/Support/ToolUtilities.cpp

49

Oh wow!

// -----// IR Dump Before CSE //----- //
func @foo() {
  return
}

// -----// IR Dump Before CSE //----- //
func @baz() {
  return
}

// -----// IR Dump Before Canonicalizer //----- //
func @foo() {
  return
}

// -----// IR Dump Before Canonicalizer //----- //
func @baz() {
  return
}

I didn't notice this trick about the extra // before, nifty!

[jpienaar]

In D109636#3009906, @rriddle wrote:

What is the cost of doing these checks on large files?

It is more a function of number of splits than file size: on files you would check ~4 characters 2x instead of 1x per potential split. You have 1 additional loop over splits, and in case where there is no near misses you have 4N characters additionally checked and stringref's inserted into another vector (which I could reserve space for all to avoid any allocations in the loop), so it is linear in number of splits. If you have near misses throughout the file you end up with 4N additional checks + N-1 StringRef creations and getFromPointer calls - so that could be heavy, but also unexpected that we'd have such near misses throughout the file.

mlir/lib/Support/ToolUtilities.cpp

49

I'll revert the newline check then (or I could additionally allow for "" or " " suffices, but not sure that valuable).

Although what is the use case? E.g., if you dump all the sections, would you feed this output file into mlir-opt to run a pass over all stages of a pipeline? (I'd have thought you'd just copy & past the part of interest out into a new file post print-ir-* rather than feed in the whole file).

[rriddle]

mlir/lib/Support/ToolUtilities.cpp

49

The main concrete use case is that you can open the dump in vscode(or another IDE) and get language server capabilities (without having the different dumps negatively interact with each other).

[jpienaar]

Removed newline check as that is desired for print-ir-* behavior

[rriddle]

mlir/lib/Support/ToolUtilities.cpp

54–55

Will this crash if the split is the last line?

[jpienaar]

Ensure no check even if buffer is not null-terminated

[jpienaar]

mlir/lib/Support/ToolUtilities.cpp

54–55

Good question, it didn't but I think that was as mostly null-terminated main buffer here, but checking explicitly better.

mlir/lib/Support/ToolUtilities.cpp

 //===- ToolUtilities.cpp - MLIR Tool Utilities ----------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines common utilities for implementing MLIR tools.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Support/ToolUtilities.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 
 LogicalResult
 mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer> originalBuffer,
                             ChunkBufferHandler processChunkBuffer,
                             raw_ostream &os) {
   const char splitMarker[] = "// -----";
+  const int kSplitMarkerLen = strlen(splitMarker);
 
   auto *origMemBuffer = originalBuffer.get();
-  SmallVector<StringRef, 8> sourceBuffers;
-  origMemBuffer->getBuffer().split(sourceBuffers, splitMarker);
+  SmallVector<StringRef, 8> rawSourceBuffers;
+  // Split dropping the last '-' to enable flagging near misses.
+  origMemBuffer->getBuffer().split(rawSourceBuffers,
+                                   StringRef(splitMarker).drop_back());
+  if (rawSourceBuffers.empty())
+    return success();
 
   // Add the original buffer to the source manager.
   llvm::SourceMgr fileSourceMgr;
   fileSourceMgr.AddNewSourceBuffer(std::move(originalBuffer), llvm::SMLoc());
 
+  // Flag near misses.
+  SmallVector<StringRef, 8> sourceBuffers;
+  StringRef prev;
+  for (auto buffer : rawSourceBuffers) {
+    if (prev.empty()) {
+      prev = buffer;
+      continue;
+    }
+
+    if (buffer.startswith("-")) {
+      sourceBuffers.push_back(prev);
+      prev = buffer.drop_front();
+    } else {
+      // TODO: Consider making this a failure.
+      auto splitLoc = llvm::SMLoc::getFromPointer(buffer.data());
+      fileSourceMgr.PrintMessage(llvm::errs(), splitLoc,
+                                 llvm::SourceMgr::DK_Warning,
+                                 "near miss with file split marker");
+      prev = StringRef(prev.data(),
+                       prev.size() + kSplitMarkerLen - 1 + buffer.size());
+    }
+  }
+  if (!prev.empty())
+    sourceBuffers.push_back(prev);
+
   // Process each chunk in turn.
   bool hadFailure = false;
   for (auto &subBuffer : sourceBuffers) {
     auto splitLoc = llvm::SMLoc::getFromPointer(subBuffer.data());
     unsigned splitLine = fileSourceMgr.getLineAndColumn(splitLoc).first;
     auto subMemBuffer = llvm::MemoryBuffer::getMemBufferCopy(
         subBuffer, Twine("within split at ") +
                        origMemBuffer->getBufferIdentifier() + ":" +
                        Twine(splitLine) + " offset ");
     if (failed(processChunkBuffer(std::move(subMemBuffer), os)))
       hadFailure = true;
   }
 
   // If any fails, then return a failure of the tool.
   return failure(hadFailure);
 }

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 // CHECK:           } : tensor<?xindex>
 // CHECK:           return
 // CHECK:         }
   %0 = shape.broadcast %a, %b, %c
       : tensor<2xindex>, tensor<3xindex>, tensor<2xindex> -> tensor<?xindex>
   return
 }
 
-// ----
+// -----
 
 // CHECK-LABEL: @broadcast_to_known_rank
 func @broadcast_to_known_rank(%a : tensor<1xindex>, %b : tensor<3xindex>)
     -> tensor<3xindex> {
   // CHECK: %[[RES:.*]] = tensor.cast %{{.*}} : tensor<?xindex> to tensor<3xindex>
   // CHECK: return %[[RES]] : tensor<3xindex>
   %0 = shape.broadcast %a, %b : tensor<1xindex>, tensor<3xindex> -> tensor<3xindex>
   return %0 : tensor<3xindex>

mlir/test/Conversion/TosaToSCF/tosa-to-scf.mlir

^bb0(%arg2: tensor<i32>):
     %3 = "tosa.add"(%arg2, %2) : (tensor<i32>, tensor<i32>) -> tensor<i32>
 
     // CHECK: scf.yield [[ADD]]
     "tosa.yield"(%3) : (tensor<i32>) -> ()
   }) : (tensor<i32>) -> (tensor<i32>)
   return %1 : tensor<i32>
 }
 
-// ----
+// -----
 
 // CHECK-LABEL: func @if_test
 // CHECK-SAME: ([[ARG0:%.+]]: tensor<f32>, [[ARG1:%.+]]: tensor<f32>, [[ARG2:%.+]]: tensor<i1>)
 func @if_test(%arg0 : tensor<f32>, %arg1 : tensor<f32>, %arg2 : tensor<i1>) -> (tensor<f32>) {
   // CHECK: [[EX:%.+]] = tensor.extract [[ARG2]]
   // CHECK: [[IF:%.+]] = scf.if [[EX]] -> (tensor<f32>) {
   %0 = "tosa.cond_if"(%arg2, %arg0, %arg1) ({
 

mlir/test/Dialect/OpenACC/canonicalize.mlir

func @testupdateop(%a: memref<10xf32>) -> () {
   %ifCond = constant false
   acc.update if(%ifCond) host(%a: memref<10xf32>)
   return
 }
 
 // CHECK: func @testupdateop
 // CHECK-NOT: acc.update
 
-// ----
+// -----
 
 func @testenterdataop(%a: memref<10xf32>, %ifCond: i1) -> () {
   acc.enter_data if(%ifCond) create(%a: memref<10xf32>)
   return
 }
 
 // CHECK:  func @testenterdataop(%{{.*}}: memref<10xf32>, [[IFCOND:%.*]]: i1)
 // CHECK:    acc.enter_data if(%{{.*}}) create(%{{.*}} : memref<10xf32>)

mlir/test/Dialect/SCF/canonicalize.mlir

func @empty_if2(%cond: i1) {
   }
   return
 }
 
 // CHECK-LABEL:   func @empty_if2
 // CHECK-NOT:       scf.if
 // CHECK:           return
 
-// ----
+// -----
 
 func @empty_else(%cond: i1, %v : memref<i1>) {
   scf.if %cond {
     memref.store %cond, %v[] : memref<i1>
   } else {
   }
   return
 }

mlir/test/Dialect/Shape/canonicalize.mlir

 // CHECK-LABEL: @cast_extent_tensor
 func @cast_extent_tensor(%arg : tensor<*xf32>) -> tensor<3xindex> {
   // CHECK: tensor.cast %{{.*}} : tensor<?xindex> to tensor<3xindex>
   %0 = shape.shape_of %arg : tensor<*xf32> -> tensor<?xindex>
   %1 = tensor.cast %0 : tensor<?xindex> to tensor<3xindex>
   return %1 : tensor<3xindex>
 }
 
-// ----
+// -----
 
 // CHECK-LABEL: max_same_arg
 // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape)
 func @max_same_arg(%a: !shape.shape) -> !shape.shape {
   %1 = shape.max %a, %a : !shape.shape, !shape.shape -> !shape.shape
   // CHECK: return %[[SHAPE]]
   return %1 : !shape.shape
 }
 
-// ----
+// -----
 
 // CHECK-LABEL: min_same_arg
 // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape)
 func @min_same_arg(%a: !shape.shape) -> !shape.shape {
   %1 = shape.min %a, %a : !shape.shape, !shape.shape -> !shape.shape
   // CHECK: return %[[SHAPE]]
   return %1 : !shape.shape
 }
-// ----
+// -----
 
 // CHECK-LABEL: @cstr_broadcastable_folding
 func @cstr_broadcastable_folding(%arg : tensor<?x4xf32>) {
   // CHECK: const_witness true
   %0 = shape.shape_of %arg : tensor<?x4xf32> -> tensor<2xindex>
   %1 = shape.const_shape [4] : tensor<1xindex>
   %2 = shape.cstr_broadcastable %0, %1: tensor<2xindex>, tensor<1xindex>
   "use"(%2) : (!shape.witness) -> ()

mlir/test/mlir-opt/nearmiss.mlir

+// RUN: mlir-opt --split-input-file --verify-diagnostics %s 2> %t &&  FileCheck --input-file %t %s
+// RUN: cat %t
+
+func @main() {return}
+
+// -----
+
+// expected-note @+1 {{see existing symbol definition here}}
+func @foo() { return }
+// CHECK: warning: near miss with file split marker
+// CHECK: ----
+// ----
+
+// expected-error @+1 {{redefinition of symbol named 'foo'}}
+func @foo() { return }
+// CHECK: warning: near miss with file split marker
+// CHECK: ----
+// ----
+func @bar2() {return }
+