[mlir] Add FileRange location type. (#80213)

jpienaar · River707 · web-flow · commit 01e75646a5d4 · 2024-11-23T05:12:04.000-08:00
This location type represents a contiguous range inside a file. It is
effectively a pair of FileLineCols. Add new type and make FileLineCol a
view for case where it matches existing previous one.

The location includes filename and optional start line &amp; col, and end
line &amp; col. Considered common cases are file:line, file:line:col,
file:line:start_col to file:line:end_col and general range within same
file. In memory its encoded as trailing objects. This keeps the memory
requirement the same as FileLineColLoc today (makes the rather common
File:Line cheaper) at the expense of extra work at decoding time. Kept the unsigned
type.

There was the option to always have file range be castable to
FileLineColLoc. This cast would just drop other fields. That may result
in some simpler staging. TBD.

This is a rather minimal change, it does not yet add bindings (C or
Python), lowering to LLVM debug locations etc. that supports end line:cols.

---------

Co-authored-by: River Riddle &lt;riddleriver@gmail.com&gt;
diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
@@ -19,6 +19,7 @@ class AffineExpr;
 class IRMapping;
 class UnknownLoc;
 class FileLineColLoc;
+class FileLineColRange;
 class Type;
 class PrimitiveType;
 class IntegerType;
diff --git a/mlir/include/mlir/IR/BuiltinDialectBytecode.td b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
@@ -95,11 +95,26 @@ def CallSiteLoc : DialectAttribute<(attr
   LocationAttr:$caller
 )>;
 
+let cType = "FileLineColRange" in {
+def FileLineColRange : DialectAttribute<(attr
+  StringAttr:$filename,
+  WithBuilder<"$_args",
+    WithType<"SmallVector<uint64_t>",
+    WithParser<"succeeded(readFileLineColRangeLocs($_reader, $_var))",
+    WithPrinter<"writeFileLineColRangeLocs($_writer, $_name)">>>>:$rawLocData
+)> {
+  let cBuilder = "getFileLineColRange(context, filename, rawLocData)";
+  let printerPredicate = "!::llvm::isa<FileLineColLoc>($_val)";
+}
+
 def FileLineColLoc : DialectAttribute<(attr
   StringAttr:$filename,
-  VarInt:$line,
-  VarInt:$column
-)>;
+  VarInt:$start_line,
+  VarInt:$start_column
+)> {
+  let printerPredicate = "::llvm::isa<FileLineColLoc>($_val)";
+}
+}
 
 let cType = "FusedLoc",
     cBuilder = "cast<FusedLoc>(get<FusedLoc>(context, $_args))" in {
@@ -321,7 +336,8 @@ def BuiltinDialectAttributes : DialectAttributes<"Builtin"> {
     DenseIntOrFPElementsAttr,
     DenseStringElementsAttr,
     SparseElementsAttr,
-    DistinctAttr
+    DistinctAttr,
+    FileLineColRange,
   ];
 }
 
diff --git a/mlir/include/mlir/IR/BuiltinLocationAttributes.td b/mlir/include/mlir/IR/BuiltinLocationAttributes.td
@@ -60,46 +60,98 @@ def CallSiteLoc : Builtin_LocationAttr<"CallSiteLoc"> {
 }
 
 //===----------------------------------------------------------------------===//
-// FileLineColLoc
+// FileLineColRange
 //===----------------------------------------------------------------------===//
 
-def FileLineColLoc : Builtin_LocationAttr<"FileLineColLoc"> {
-  let summary = "A file:line:column source location";
+def FileLineColRange : Builtin_LocationAttr<"FileLineColRange"> {
+  let summary = "A file:line:column source location range";
   let description = [{
     Syntax:
 
     ```
     filelinecol-location ::= string-literal `:` integer-literal `:`
                              integer-literal
+                             (`to` (integer-literal ?) `:` integer-literal ?)
     ```
 
-    An instance of this location represents a tuple of file, line number, and
-    column number. This is similar to the type of location that you get from
-    most source languages.
+    An instance of this location represents a tuple of file, start and end line
+    number, and start and end column number. It allows for the following
+    configurations:
+
+    *   A single file line location: `file:line`;
+    *   A single file line col location: `file:line:column`;
+    *   A single line range: `file:line:column to :column`;
+    *   A single file range: `file:line:column to line:column`;
 
     Example:
 
     ```mlir
-    loc("mysource.cc":10:8)
+    loc("mysource.cc":10:8 to 12:18)
     ```
   }];
-  let parameters = (ins "StringAttr":$filename, "unsigned":$line,
-                        "unsigned":$column);
+
+  let parameters = (ins "StringAttr":$filename,
+    "unsigned":$start_line, "unsigned":$start_column,
+    "unsigned":$end_line, "unsigned":$end_column);
   let builders = [
+    AttrBuilderWithInferredContext<(ins "StringAttr":$filename), [{
+      return $_get(filename.getContext(), filename, ArrayRef<unsigned>{});
+    }]>,
+    AttrBuilderWithInferredContext<(ins "StringAttr":$filename,
+                                        "unsigned":$line), [{
+      return $_get(filename.getContext(), filename,
+                   ArrayRef<unsigned>{line});
+    }]>,
     AttrBuilderWithInferredContext<(ins "StringAttr":$filename,
                                         "unsigned":$line,
                                         "unsigned":$column), [{
-      return $_get(filename.getContext(), filename, line, column);
+      return $_get(filename.getContext(), filename,
+                   ArrayRef<unsigned>{line, column});
     }]>,
-    AttrBuilder<(ins "StringRef":$filename, "unsigned":$line,
-                     "unsigned":$column), [{
+    AttrBuilder<(ins "::llvm::StringRef":$filename,
+                     "unsigned":$start_line,
+                     "unsigned":$start_column), [{
       return $_get($_ctxt,
-                   StringAttr::get($_ctxt, filename.empty() ? "-" : filename),
-                   line, column);
-    }]>
+        StringAttr::get($_ctxt, filename.empty() ? "-" : filename),
+        ArrayRef<unsigned>{start_line, start_column});
+    }]>,
+    AttrBuilderWithInferredContext<(ins "::mlir::StringAttr":$filename,
+                                        "unsigned":$line,
+                                        "unsigned":$start_column,
+                                        "unsigned":$end_column), [{
+      return $_get(filename.getContext(), filename,
+                   ArrayRef<unsigned>{line, start_column, end_column});
+    }]>,
+    AttrBuilderWithInferredContext<(ins "::mlir::StringAttr":$filename,
+                                        "unsigned":$start_line,
+                                        "unsigned":$start_column,
+                                        "unsigned":$end_line,
+                                        "unsigned":$end_column), [{
+      return $_get(filename.getContext(), filename,
+        ArrayRef<unsigned>{start_line, start_column, end_column, end_line});
+    }]>,
+    AttrBuilder<(ins "::llvm::StringRef":$filename,
+                     "unsigned":$start_line,
+                     "unsigned":$start_column,
+                     "unsigned":$end_line,
+                     "unsigned":$end_column), [{
+      return $_get($_ctxt,
+        StringAttr::get($_ctxt, filename.empty() ? "-" : filename),
+        ArrayRef<unsigned>{start_line, start_column, end_column, end_line});
+    }]>,
   ];
+
+  let extraClassDeclaration = [{
+    ::mlir::StringAttr getFilename() const;
+    unsigned getStartLine() const;
+    unsigned getStartColumn() const;
+    unsigned getEndColumn() const;
+    unsigned getEndLine() const;
+  }];
   let skipDefaultBuilders = 1;
-  let attrName = "builtin.file_line_loc";
+  let genAccessors = 0;
+  let genStorageClass = 0;
+  let attrName = "builtin.file_line_range";
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/IR/Location.h b/mlir/include/mlir/IR/Location.h
@@ -136,6 +136,11 @@ inline ::llvm::hash_code hash_value(Location arg) {
 // Tablegen Attribute Declarations
 //===----------------------------------------------------------------------===//
 
+// Forward declaration for class created later.
+namespace mlir::detail {
+struct FileLineColRangeAttrStorage;
+} // namespace mlir::detail
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/IR/BuiltinLocationAttributes.h.inc"
 
@@ -164,6 +169,32 @@ class FusedLocWith : public FusedLoc {
   }
 };
 
+//===----------------------------------------------------------------------===//
+// FileLineColLoc
+//===----------------------------------------------------------------------===//
+
+/// An instance of this location represents a tuple of file, line number, and
+/// column number. This is similar to the type of location that you get from
+/// most source languages.
+///
+/// FileLineColLoc is a FileLineColRange with exactly one line and column.
+class FileLineColLoc : public FileLineColRange {
+public:
+  using FileLineColRange::FileLineColRange;
+
+  static FileLineColLoc get(StringAttr filename, unsigned line,
+                            unsigned column);
+  static FileLineColLoc get(MLIRContext *context, StringRef fileName,
+                            unsigned line, unsigned column);
+
+  StringAttr getFilename() const;
+  unsigned getLine() const;
+  unsigned getColumn() const;
+
+  /// Methods for support type inquiry through isa, cast, and dyn_cast.
+  static bool classof(Attribute attr);
+};
+
 //===----------------------------------------------------------------------===//
 // OpaqueLoc
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/AsmParser/LocationParser.cpp b/mlir/lib/AsmParser/LocationParser.cpp
@@ -12,6 +12,7 @@
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Location.h"
 #include "mlir/Support/LLVM.h"
+#include <optional>
 
 using namespace mlir;
 using namespace mlir::detail;
@@ -97,37 +98,82 @@ ParseResult Parser::parseFusedLocation(LocationAttr &loc) {
   return success();
 }
 
-ParseResult Parser::parseNameOrFileLineColLocation(LocationAttr &loc) {
+ParseResult Parser::parseNameOrFileLineColRange(LocationAttr &loc) {
   auto *ctx = getContext();
   auto str = getToken().getStringValue();
   consumeToken(Token::string);
 
+  std::optional<unsigned> startLine, startColumn, endLine, endColumn;
+
   // If the next token is ':' this is a filelinecol location.
   if (consumeIf(Token::colon)) {
     // Parse the line number.
     if (getToken().isNot(Token::integer))
       return emitWrongTokenError(
-          "expected integer line number in FileLineColLoc");
-    auto line = getToken().getUnsignedIntegerValue();
-    if (!line)
+          "expected integer line number in FileLineColRange");
+    startLine = getToken().getUnsignedIntegerValue();
+    if (!startLine)
       return emitWrongTokenError(
-          "expected integer line number in FileLineColLoc");
+          "expected integer line number in FileLineColRange");
     consumeToken(Token::integer);
 
     // Parse the ':'.
-    if (parseToken(Token::colon, "expected ':' in FileLineColLoc"))
-      return failure();
+    if (getToken().isNot(Token::colon)) {
+      loc = FileLineColRange::get(StringAttr::get(ctx, str), *startLine);
+      return success();
+    }
+    consumeToken(Token::colon);
 
     // Parse the column number.
-    if (getToken().isNot(Token::integer))
+    if (getToken().isNot(Token::integer)) {
+      return emitWrongTokenError(
+          "expected integer column number in FileLineColRange");
+    }
+    startColumn = getToken().getUnsignedIntegerValue();
+    if (!startColumn.has_value())
+      return emitError("expected integer column number in FileLineColRange");
+    consumeToken(Token::integer);
+
+    if (!isCurrentTokenAKeyword() || getTokenSpelling() != "to") {
+      loc = FileLineColLoc::get(ctx, str, *startLine, *startColumn);
+      return success();
+    }
+    consumeToken();
+
+    // Parse the line number.
+    if (getToken().is(Token::integer)) {
+      endLine = getToken().getUnsignedIntegerValue();
+      if (!endLine) {
+        return emitWrongTokenError(
+            "expected integer line number in FileLineColRange");
+      }
+      consumeToken(Token::integer);
+    }
+
+    // Parse the ':'.
+    if (getToken().isNot(Token::colon)) {
+      return emitWrongTokenError(
+          "expected either integer or `:` post `to` in FileLineColRange");
+    }
+    consumeToken(Token::colon);
+
+    // Parse the column number.
+    if (getToken().isNot(Token::integer)) {
       return emitWrongTokenError(
-          "expected integer column number in FileLineColLoc");
-    auto column = getToken().getUnsignedIntegerValue();
-    if (!column.has_value())
-      return emitError("expected integer column number in FileLineColLoc");
+          "expected integer column number in FileLineColRange");
+    }
+    endColumn = getToken().getUnsignedIntegerValue();
+    if (!endColumn.has_value())
+      return emitError("expected integer column number in FileLineColRange");
     consumeToken(Token::integer);
 
-    loc = FileLineColLoc::get(ctx, str, *line, *column);
+    if (endLine.has_value()) {
+      loc = FileLineColRange::get(StringAttr::get(ctx, str), *startLine,
+                                  *startColumn, *endLine, *endColumn);
+    } else {
+      loc = FileLineColRange::get(StringAttr::get(ctx, str), *startLine,
+                                  *startColumn, *endColumn);
+    }
     return success();
   }
 
@@ -166,7 +212,7 @@ ParseResult Parser::parseLocationInstance(LocationAttr &loc) {
 
   // Handle either name or filelinecol locations.
   if (getToken().is(Token::string))
-    return parseNameOrFileLineColLocation(loc);
+    return parseNameOrFileLineColRange(loc);
 
   // Bare tokens required for other cases.
   if (!getToken().is(Token::bare_identifier))
diff --git a/mlir/lib/AsmParser/Parser.h b/mlir/lib/AsmParser/Parser.h
@@ -310,7 +310,7 @@ class Parser {
   ParseResult parseFusedLocation(LocationAttr &loc);
 
   /// Parse a name or FileLineCol location instance.
-  ParseResult parseNameOrFileLineColLocation(LocationAttr &loc);
+  ParseResult parseNameOrFileLineColRange(LocationAttr &loc);
 
   //===--------------------------------------------------------------------===//
   // Affine Parsing
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
@@ -2009,12 +2009,23 @@ void AsmPrinter::Impl::printLocationInternal(LocationAttr loc, bool pretty,
         else
           os << "unknown";
       })
-      .Case<FileLineColLoc>([&](FileLineColLoc loc) {
+      .Case<FileLineColRange>([&](FileLineColRange loc) {
         if (pretty)
           os << loc.getFilename().getValue();
         else
           printEscapedString(loc.getFilename());
-        os << ':' << loc.getLine() << ':' << loc.getColumn();
+        if (loc.getEndColumn() == loc.getStartColumn() &&
+            loc.getStartLine() == loc.getEndLine()) {
+          os << ':' << loc.getStartLine() << ':' << loc.getStartColumn();
+          return;
+        }
+        if (loc.getStartLine() == loc.getEndLine()) {
+          os << ':' << loc.getStartLine() << ':' << loc.getStartColumn()
+             << " to :" << loc.getEndColumn();
+          return;
+        }
+        os << ':' << loc.getStartLine() << ':' << loc.getStartColumn() << " to "
+           << loc.getEndLine() << ':' << loc.getEndColumn();
       })
       .Case<NameLoc>([&](NameLoc loc) {
         printEscapedString(loc.getName());
diff --git a/mlir/lib/IR/BuiltinDialectBytecode.cpp b/mlir/lib/IR/BuiltinDialectBytecode.cpp
diff --git a/mlir/lib/IR/Location.cpp b/mlir/lib/IR/Location.cpp
diff --git a/mlir/test/IR/locations.mlir b/mlir/test/IR/locations.mlir
