[mlir] Allow to use vformat utility with MLIR classes

[vinograd47]

Summary

Make raw_ostream operator<< follow const correctness semantic,
since it is a requirement of FormatVariadic implementation.

[rriddle]

mlir/include/mlir/IR/OpDefinition.h

112–118

I'd rather move the const_cast to a raw_ostream specific method, so that const doesn't seep into the core classes. Is that possible?

[vinograd47]

Applied review remark and reverted unnecessary format changes

[mehdi_amini]

mlir/unittests/IR/OperationSupportTest.cpp

222

This is leaky: https://lab.llvm.org/staging/#/builders/191/builds/350

Can you make sure to destroy it?

[vinograd47]

mlir/unittests/IR/OperationSupportTest.cpp

222

Done in https://github.com/llvm/llvm-project/commit/17664a256f5f928578ef13d947f2bacf3adda021

mlir/include/mlir/Dialect/MemRef/IR/MemRef.h

 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 namespace mlir {
 
 class Location;
 class OpBuilder;
 
-raw_ostream &operator<<(raw_ostream &os, Range &range);
+raw_ostream &operator<<(raw_ostream &os, const Range &range);
 
 /// Return the list of Range (i.e. offset, size, stride). Each Range
 /// entry contains either the dynamic value or a ConstantIndexOp constructed
 /// with `b` at location `loc`.
 SmallVector<Range, 8> getOrCreateRanges(OffsetSizeAndStrideOpInterface op,
                                         OpBuilder &b, Location loc);
 } // namespace mlir
 

mlir/include/mlir/IR/OpDefinition.h

public:
 
   /// Return the operation that this refers to.
   Operation *getOperation() { return state; }
 
   /// Return the context this operation belongs to.
   MLIRContext *getContext() { return getOperation()->getContext(); }
 
   /// Print the operation to the given stream.
   void print(raw_ostream &os, OpPrintingFlags flags = llvm::None) {
     state->print(os, flags);
   }
   void print(raw_ostream &os, AsmState &asmState,
              OpPrintingFlags flags = llvm::None) {
     state->print(os, asmState, flags);
   }
 
   /// Dump this operation.
   void dump() { state->dump(); }
 
   /// The source location the operation was defined or derived from.
   Location getLoc() { return state->getLoc(); }
 
   /// Return true if there are no users of any results of this operation.
inline raw_ostream &operator<<(raw_ostream &os, OpFoldResult ofr) {
   if (Value value = ofr.dyn_cast<Value>())
     value.print(os);
   else
     ofr.dyn_cast<Attribute>().print(os);
   return os;
 }
 
 /// Allow printing to a stream.
-inline raw_ostream &operator<<(raw_ostream &os, OpState &op) {
+inline raw_ostream &operator<<(raw_ostream &os, OpState op) {
   op.print(os, OpPrintingFlags().useLocalScope());
   return os;
 }
 
 //===----------------------------------------------------------------------===//
 // Operation Trait Types
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/IR/Operation.h

private:
   friend llvm::TrailingObjects<Operation, BlockOperand, Region,
                                detail::OperandStorage>;
   size_t numTrailingObjects(OverloadToken<BlockOperand>) const {
     return numSuccs;
   }
   size_t numTrailingObjects(OverloadToken<Region>) const { return numRegions; }
 };
 
-inline raw_ostream &operator<<(raw_ostream &os, Operation &op) {
-  op.print(os, OpPrintingFlags().useLocalScope());
+inline raw_ostream &operator<<(raw_ostream &os, const Operation &op) {
+  const_cast<Operation &>(op).print(os, OpPrintingFlags().useLocalScope());
   return os;
 }
 
 } // end namespace mlir
 
 namespace llvm {
 /// Provide isa functionality for operation casts.
 template <typename T> struct isa_impl<T, ::mlir::Operation> {

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

auto expectedType = SubViewOp::inferResultType(
       extractFromI64ArrayAttr(op.static_strides()));
 
   std::string errMsg;
   auto result =
       isRankReducedType(expectedType, subViewType, op.static_sizes(), &errMsg);
   return produceSubViewErrorMsg(result, op, expectedType, errMsg);
 }
 
-raw_ostream &mlir::operator<<(raw_ostream &os, Range &range) {
+raw_ostream &mlir::operator<<(raw_ostream &os, const Range &range) {
   return os << "range " << range.offset << ":" << range.size << ":"
             << range.stride;
 }
 
 /// Return the list of Range (i.e. offset, size, stride). Each Range
 /// entry contains either the dynamic value or a ConstantIndexOp constructed
 /// with `b` at location `loc`.
 SmallVector<Range, 8> mlir::getOrCreateRanges(OffsetSizeAndStrideOpInterface op,

mlir/unittests/IR/OperationSupportTest.cpp

 //===- OperationSupportTest.cpp - Operation support unit tests ------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "llvm/ADT/BitVector.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "gtest/gtest.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 static Operation *createOp(MLIRContext *context,
                            ArrayRef<Value> operands = llvm::None,
                            ArrayRef<Type> resultTypes = llvm::None,
for (int i = 0; i < kNumOpsToInsert; ++i) {
     ASSERT_TRUE(op->isBeforeInBlock(backOp));
     // Note verifyOpOrder() returns false if the order is valid.
     ASSERT_FALSE(block->verifyOpOrder());
   }
 
   containerOp->destroy();
 }
 
+TEST(OperationFormatPrintTest, CanUseVariadicFormat) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  Operation *op = createOp(&context);
+
+  std::string str = formatv("{0}", *op).str();
+  ASSERT_STREQ(str.c_str(), "\"foo.bar\"() : () -> ()");
+}
+
 } // end namespace