Add explicit friendship with Operation to each derived op class to ensure access...

--

PiperOrigin-RevId: 246661686

diff --git a/mlir/examples/Linalg/Linalg1/include/linalg1/RangeOp.h b/mlir/examples/Linalg/Linalg1/include/linalg1/RangeOp.h
index 9652f51..1973b46 100644 (file)
--- a/mlir/examples/Linalg/Linalg1/include/linalg1/RangeOp.h
+++ b/mlir/examples/Linalg/Linalg1/include/linalg1/RangeOp.h
@@ -31,6 +31,7 @@ class RangeOp : public mlir::Op<RangeOp, mlir::OpTrait::NOperands<3>::Impl,
                                 mlir::OpTrait::OneResult,
                                 mlir::OpTrait::HasNoSideEffect> {
 public:
+  friend mlir::Operation;
   using Op::Op;
 
   //////////////////////////////////////////////////////////////////////////////

diff --git a/mlir/examples/Linalg/Linalg1/include/linalg1/SliceOp.h b/mlir/examples/Linalg/Linalg1/include/linalg1/SliceOp.h
index 1d79784..48f119f 100644 (file)
--- a/mlir/examples/Linalg/Linalg1/include/linalg1/SliceOp.h
+++ b/mlir/examples/Linalg/Linalg1/include/linalg1/SliceOp.h
@@ -30,6 +30,7 @@ class SliceOp : public mlir::Op<SliceOp, mlir::OpTrait::NOperands<2>::Impl,
                                 mlir::OpTrait::OneResult,
                                 mlir::OpTrait::HasNoSideEffect> {
 public:
+  friend mlir::Operation;
   using Op::Op;
 
   //////////////////////////////////////////////////////////////////////////////

diff --git a/mlir/examples/Linalg/Linalg1/include/linalg1/ViewOp.h b/mlir/examples/Linalg/Linalg1/include/linalg1/ViewOp.h
index fcda553..54af9a6 100644 (file)
--- a/mlir/examples/Linalg/Linalg1/include/linalg1/ViewOp.h
+++ b/mlir/examples/Linalg/Linalg1/include/linalg1/ViewOp.h
@@ -32,6 +32,7 @@ class ViewOp : public mlir::Op<ViewOp, mlir::OpTrait::VariadicOperands,
                                mlir::OpTrait::OneResult,
                                mlir::OpTrait::HasNoSideEffect> {
 public:
+  friend mlir::Operation;
   using Op::Op;
 
   //////////////////////////////////////////////////////////////////////////////

diff --git a/mlir/examples/Linalg/Linalg2/include/linalg2/TensorOps.h b/mlir/examples/Linalg/Linalg2/include/linalg2/TensorOps.h
index 39e51f0..5eccbe6 100644 (file)
--- a/mlir/examples/Linalg/Linalg2/include/linalg2/TensorOps.h
+++ b/mlir/examples/Linalg/Linalg2/include/linalg2/TensorOps.h
@@ -102,6 +102,7 @@ class DotOp : public TensorContractionBase<DotOp>,
               public mlir::Op<DotOp, mlir::OpTrait::VariadicOperands,
                               mlir::OpTrait::ZeroResult> {
 public:
+  friend mlir::Operation;
   using Op::Op;
   using TensorContractionBaseType =
       TensorContractionBase::TensorContractionBaseType;
@@ -162,6 +163,7 @@ class MatvecOp : public TensorContractionBase<MatvecOp>,
                  public mlir::Op<MatvecOp, mlir::OpTrait::VariadicOperands,
                                  mlir::OpTrait::ZeroResult> {
 public:
+  friend mlir::Operation;
   using Op::Op;
   using TensorContractionBaseType =
       TensorContractionBase::TensorContractionBaseType;
@@ -222,6 +224,7 @@ class MatmulOp : public TensorContractionBase<MatmulOp>,
                  public mlir::Op<MatmulOp, mlir::OpTrait::VariadicOperands,
                                  mlir::OpTrait::ZeroResult> {
 public:
+  friend mlir::Operation;
   using Op::Op;
   using TensorContractionBaseType =
       TensorContractionBase::TensorContractionBaseType;

diff --git a/mlir/examples/Linalg/Linalg3/include/linalg3/LoadStoreOps.h b/mlir/examples/Linalg/Linalg3/include/linalg3/LoadStoreOps.h
index b77e702..a227613 100644 (file)
--- a/mlir/examples/Linalg/Linalg3/include/linalg3/LoadStoreOps.h
+++ b/mlir/examples/Linalg/Linalg3/include/linalg3/LoadStoreOps.h
@@ -30,6 +30,7 @@ class ViewType;
 class LoadOp : public mlir::Op<LoadOp, mlir::OpTrait::VariadicOperands,
                                mlir::OpTrait::OneResult> {
 public:
+  friend mlir::Operation;
   using Op::Op;
 
   //////////////////////////////////////////////////////////////////////////////
@@ -59,6 +60,7 @@ public:
 class StoreOp : public mlir::Op<StoreOp, mlir::OpTrait::VariadicOperands,
                                 mlir::OpTrait::ZeroResult> {
 public:
+  friend mlir::Operation;
   using Op::Op;
 
   //////////////////////////////////////////////////////////////////////////////

diff --git a/mlir/examples/toy/Ch3/include/toy/Dialect.h b/mlir/examples/toy/Ch3/include/toy/Dialect.h
index 84bd98e..0ad95d4 100644 (file)
--- a/mlir/examples/toy/Ch3/include/toy/Dialect.h
+++ b/mlir/examples/toy/Ch3/include/toy/Dialect.h
@@ -145,6 +145,7 @@ public:
                     mlir::FloatAttr value);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -181,6 +182,7 @@ public:
   llvm::StringRef getCalleeName();
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -210,6 +212,7 @@ public:
   mlir::Value *getOperand() { return getOperation()->getOperand(0); }
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -230,6 +233,7 @@ public:
                     mlir::Value *value);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -249,6 +253,7 @@ public:
                     mlir::Value *value);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -273,6 +278,7 @@ public:
                     mlir::Value *value, ToyArrayType reshapedType);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -301,6 +307,7 @@ public:
   mlir::Value *getRHS() { return getOperand(1); }
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -327,6 +334,7 @@ public:
   mlir::Value *getRHS() { return getOperand(1); }
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 

diff --git a/mlir/examples/toy/Ch4/include/toy/Dialect.h b/mlir/examples/toy/Ch4/include/toy/Dialect.h
index 96e05dc..b0a57d9 100644 (file)
--- a/mlir/examples/toy/Ch4/include/toy/Dialect.h
+++ b/mlir/examples/toy/Ch4/include/toy/Dialect.h
@@ -145,6 +145,7 @@ public:
                     mlir::FloatAttr value);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -181,6 +182,7 @@ public:
   llvm::StringRef getCalleeName();
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -210,6 +212,7 @@ public:
   mlir::Value *getOperand() { return getOperation()->getOperand(0); }
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -230,6 +233,7 @@ public:
                     mlir::Value *value);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -255,6 +259,7 @@ public:
                               mlir::MLIRContext *context);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -285,6 +290,7 @@ public:
                               mlir::MLIRContext *context);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -314,6 +320,7 @@ public:
   mlir::Value *getRHS() { return getOperand(1); }
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -341,6 +348,7 @@ public:
   mlir::Value *getRHS() { return getOperand(1); }
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 

diff --git a/mlir/examples/toy/Ch5/include/toy/Dialect.h b/mlir/examples/toy/Ch5/include/toy/Dialect.h
index 9d7f82d..6b6fe79 100644 (file)
--- a/mlir/examples/toy/Ch5/include/toy/Dialect.h
+++ b/mlir/examples/toy/Ch5/include/toy/Dialect.h
@@ -153,6 +153,7 @@ public:
   }
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -189,6 +190,7 @@ public:
   llvm::StringRef getCalleeName();
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -218,6 +220,7 @@ public:
   mlir::Value *getOperand() { return getOperation()->getOperand(0); }
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -238,6 +241,7 @@ public:
                     mlir::Value *value);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -263,6 +267,7 @@ public:
                               mlir::MLIRContext *context);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -293,6 +298,7 @@ public:
                               mlir::MLIRContext *context);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -322,6 +328,7 @@ public:
   mlir::Value *getRHS() { return getOperand(1); }
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -349,6 +356,7 @@ public:
   mlir::Value *getRHS() { return getOperand(1); }
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -366,6 +374,7 @@ public:
                     mlir::Type retType);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 
@@ -385,6 +394,7 @@ public:
                               mlir::MLIRContext *context);
 
   /// Inherit constructor.
+  friend mlir::Operation;
   using Op::Op;
 };
 

diff --git a/mlir/include/mlir/AffineOps/AffineOps.h b/mlir/include/mlir/AffineOps/AffineOps.h
index 5b1446d..d46ecc5 100644 (file)
--- a/mlir/include/mlir/AffineOps/AffineOps.h
+++ b/mlir/include/mlir/AffineOps/AffineOps.h
@@ -59,6 +59,7 @@ public:
 class AffineApplyOp : public Op<AffineApplyOp, OpTrait::VariadicOperands,
                                 OpTrait::OneResult, OpTrait::HasNoSideEffect> {
 public:
+  friend Operation;
   using Op::Op;
 
   /// Builds an affine apply op with the specified map and operands.
@@ -118,6 +119,7 @@ public:
 class AffineForOp
     : public Op<AffineForOp, OpTrait::VariadicOperands, OpTrait::ZeroResult> {
 public:
+  friend Operation;
   using Op::Op;
 
   // Hooks to customize behavior of this op.
@@ -304,6 +306,7 @@ private:
 class AffineIfOp
     : public Op<AffineIfOp, OpTrait::VariadicOperands, OpTrait::ZeroResult> {
 public:
+  friend Operation;
   using Op::Op;
 
   // Hooks to customize behavior of this op.
@@ -337,6 +340,7 @@ class AffineTerminatorOp
     : public Op<AffineTerminatorOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
                 OpTrait::IsTerminator> {
 public:
+  friend Operation;
   using Op::Op;
 
   static void build(Builder *, OperationState *) {}

diff --git a/mlir/include/mlir/GPU/GPUDialect.h b/mlir/include/mlir/GPU/GPUDialect.h
index 555c644..2229b2a 100644 (file)
--- a/mlir/include/mlir/GPU/GPUDialect.h
+++ b/mlir/include/mlir/GPU/GPUDialect.h
@@ -49,6 +49,7 @@ class LaunchOp : public Op<LaunchOp, OpTrait::AtLeastNOperands<6>::Impl,
                            OpTrait::ZeroResult,
                            OpTrait::NthRegionIsIsolatedAbove<0>::Impl> {
 public:
+  friend Operation;
   using Op::Op;
 
   static void build(Builder *builder, OperationState *result, Value *gridSizeX,

diff --git a/mlir/include/mlir/Linalg/IR/LinalgOps.h b/mlir/include/mlir/Linalg/IR/LinalgOps.h
index 5503b65..58ff5a9 100644 (file)
--- a/mlir/include/mlir/Linalg/IR/LinalgOps.h
+++ b/mlir/include/mlir/Linalg/IR/LinalgOps.h
@@ -36,6 +36,7 @@ namespace mlir {
 class BufferAllocOp
     : public Op<BufferAllocOp, OpTrait::OneOperand, OpTrait::OneResult> {
 public:
+  friend Operation;
   using Op::Op;
 
   // Hooks to customize the behavior of this op.
@@ -59,6 +60,7 @@ public:
 class BufferDeallocOp
     : public Op<BufferDeallocOp, OpTrait::OneOperand, OpTrait::ZeroResult> {
 public:
+  friend Operation;
   using Op::Op;
 
   // Hooks to customize the behavior of this op.
@@ -84,6 +86,7 @@ public:
 class RangeOp : public Op<RangeOp, OpTrait::NOperands<3>::Impl,
                           OpTrait::OneResult, OpTrait::HasNoSideEffect> {
 public:
+  friend Operation;
   using Op::Op;
 
   // Hooks to customize the behavior of this op.
@@ -145,6 +148,7 @@ class SliceOp : public mlir::Op<SliceOp, mlir::OpTrait::VariadicOperands,
   enum { FirstIndexingOperand = 1 };
 
 public:
+  friend Operation;
   using Op::Op;
 
   // Hooks to customize the behavior of this op.
@@ -194,6 +198,7 @@ class ViewOp : public mlir::Op<ViewOp, mlir::OpTrait::VariadicOperands,
   enum { FirstIndexingOperand = 1 };
 
 public:
+  friend Operation;
   using Op::Op;
 
   // Hooks to customize the behavior of this op.

diff --git a/mlir/include/mlir/StandardOps/Ops.h b/mlir/include/mlir/StandardOps/Ops.h
index fbb5cde..fd2cb88 100644 (file)
--- a/mlir/include/mlir/StandardOps/Ops.h
+++ b/mlir/include/mlir/StandardOps/Ops.h
@@ -68,6 +68,7 @@ public:
 class AllocOp
     : public Op<AllocOp, OpTrait::VariadicOperands, OpTrait::OneResult> {
 public:
+  friend Operation;
   using Op::Op;
 
   /// The result of an alloc is always a MemRefType.
@@ -98,6 +99,7 @@ public:
 class BranchOp : public Op<BranchOp, OpTrait::VariadicOperands,
                            OpTrait::ZeroResult, OpTrait::IsTerminator> {
 public:
+  friend Operation;
   using Op::Op;
 
   static StringRef getOperationName() { return "std.br"; }
@@ -126,6 +128,7 @@ public:
 class CallOp
     : public Op<CallOp, OpTrait::VariadicOperands, OpTrait::VariadicResults> {
 public:
+  friend Operation;
   using Op::Op;
 
   static StringRef getOperationName() { return "std.call"; }
@@ -162,6 +165,7 @@ public:
 class CallIndirectOp : public Op<CallIndirectOp, OpTrait::VariadicOperands,
                                  OpTrait::VariadicResults> {
 public:
+  friend Operation;
   using Op::Op;
   static StringRef getOperationName() { return "std.call_indirect"; }
 
@@ -231,6 +235,7 @@ class CmpIOp
                 OpTrait::OneResult, OpTrait::ResultsAreBoolLike,
                 OpTrait::SameOperandsAndResultShape, OpTrait::HasNoSideEffect> {
 public:
+  friend Operation;
   using Op::Op;
 
   CmpIPredicate getPredicate() {
@@ -272,6 +277,7 @@ class CondBranchOp : public Op<CondBranchOp, OpTrait::AtLeastNOperands<1>::Impl,
   /// follows:
   /// { condition, [true_operands], [false_operands] }
 public:
+  friend Operation;
   using Op::Op;
 
   static StringRef getOperationName() { return "std.cond_br"; }
@@ -363,6 +369,7 @@ private:
 ///
 class ConstantFloatOp : public ConstantOp {
 public:
+  friend Operation;
   using ConstantOp::ConstantOp;
 
   /// Builds a constant float op producing a float of the specified type.
@@ -381,6 +388,7 @@ public:
 ///
 class ConstantIntOp : public ConstantOp {
 public:
+  friend Operation;
   using ConstantOp::ConstantOp;
   /// Build a constant int op producing an integer of the specified width.
   static void build(Builder *builder, OperationState *result, int64_t value,
@@ -403,6 +411,7 @@ public:
 ///
 class ConstantIndexOp : public ConstantOp {
 public:
+  friend Operation;
   using ConstantOp::ConstantOp;
 
   /// Build a constant int op producing an index.
@@ -426,6 +435,7 @@ public:
 class DeallocOp
     : public Op<DeallocOp, OpTrait::OneOperand, OpTrait::ZeroResult> {
 public:
+  friend Operation;
   using Op::Op;
 
   Value *getMemRef() { return getOperand(); }
@@ -451,6 +461,7 @@ public:
 class DimOp : public Op<DimOp, OpTrait::OneOperand, OpTrait::OneResult,
                         OpTrait::HasNoSideEffect> {
 public:
+  friend Operation;
   using Op::Op;
 
   static void build(Builder *builder, OperationState *result,
@@ -511,6 +522,7 @@ public:
 class DmaStartOp
     : public Op<DmaStartOp, OpTrait::VariadicOperands, OpTrait::ZeroResult> {
 public:
+  friend Operation;
   using Op::Op;
 
   static void build(Builder *builder, OperationState *result, Value *srcMemRef,
@@ -636,6 +648,7 @@ public:
 class DmaWaitOp
     : public Op<DmaWaitOp, OpTrait::VariadicOperands, OpTrait::ZeroResult> {
 public:
+  friend Operation;
   using Op::Op;
 
   static void build(Builder *builder, OperationState *result, Value *tagMemRef,
@@ -681,6 +694,7 @@ class ExtractElementOp
     : public Op<ExtractElementOp, OpTrait::VariadicOperands, OpTrait::OneResult,
                 OpTrait::HasNoSideEffect> {
 public:
+  friend Operation;
   using Op::Op;
 
   static void build(Builder *builder, OperationState *result, Value *aggregate,
@@ -712,6 +726,7 @@ public:
 class LoadOp
     : public Op<LoadOp, OpTrait::VariadicOperands, OpTrait::OneResult> {
 public:
+  friend Operation;
   using Op::Op;
 
   // Hooks to customize behavior of this op.
@@ -754,6 +769,7 @@ public:
 ///
 class MemRefCastOp : public CastOp<MemRefCastOp> {
 public:
+  friend Operation;
   using CastOp::CastOp;
   static StringRef getOperationName() { return "std.memref_cast"; }
 
@@ -781,6 +797,7 @@ public:
 class ReturnOp : public Op<ReturnOp, OpTrait::VariadicOperands,
                            OpTrait::ZeroResult, OpTrait::IsTerminator> {
 public:
+  friend Operation;
   using Op::Op;
 
   static StringRef getOperationName() { return "std.return"; }
@@ -808,6 +825,7 @@ public:
 class SelectOp : public Op<SelectOp, OpTrait::NOperands<3>::Impl,
                            OpTrait::OneResult, OpTrait::HasNoSideEffect> {
 public:
+  friend Operation;
   using Op::Op;
 
   static StringRef getOperationName() { return "std.select"; }
@@ -837,6 +855,7 @@ public:
 class StoreOp
     : public Op<StoreOp, OpTrait::VariadicOperands, OpTrait::ZeroResult> {
 public:
+  friend Operation;
   using Op::Op;
 
   // Hooks to customize behavior of this op.
@@ -877,6 +896,7 @@ public:
 ///
 class TensorCastOp : public CastOp<TensorCastOp> {
 public:
+  friend Operation;
   using CastOp::CastOp;
 
   static StringRef getOperationName() { return "std.tensor_cast"; }

diff --git a/mlir/include/mlir/VectorOps/VectorOps.h b/mlir/include/mlir/VectorOps/VectorOps.h
index 6b75933..30e5161 100644 (file)
--- a/mlir/include/mlir/VectorOps/VectorOps.h
+++ b/mlir/include/mlir/VectorOps/VectorOps.h
@@ -93,6 +93,7 @@ class VectorTransferReadOp
   enum Offsets : unsigned { MemRefOffset = 0, FirstIndexOffset = 1 };
 
 public:
+  friend Operation;
   using Op::Op;
 
   static StringRef getOperationName() { return "vector.transfer_read"; }
@@ -157,6 +158,7 @@ class VectorTransferWriteOp
   };
 
 public:
+  friend Operation;
   using Op::Op;
 
   static StringRef getOperationName() { return "vector.transfer_write"; }
@@ -192,6 +194,7 @@ public:
 class VectorTypeCastOp
     : public Op<VectorTypeCastOp, OpTrait::OneOperand, OpTrait::OneResult> {
 public:
+  friend Operation;
   using Op::Op;
 
   static StringRef getOperationName() { return "vector.type_cast"; }

diff --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
index 1a31673..de7bab2 100644 (file)
--- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
@@ -307,6 +307,7 @@ void OpClass::writeDeclTo(raw_ostream &os) const {
   for (const auto &trait : traits)
     os << ", " << trait;
   os << "> {\npublic:\n";
+  os << "  friend Operation;\n";
   os << "  using Op::Op;\n";
   for (const auto &method : methods) {
     method.writeDeclTo(os);