Add a general operation property 'IsolatedFromAbove' that guarantees that all regions...

PiperOrigin-RevId: 251927466

diff --git a/mlir/include/mlir/GPU/GPUDialect.h b/mlir/include/mlir/GPU/GPUDialect.h
index 003336b..2ba320b 100644 (file)
--- a/mlir/include/mlir/GPU/GPUDialect.h
+++ b/mlir/include/mlir/GPU/GPUDialect.h
@@ -56,8 +56,7 @@ struct KernelDim3 {
 /// the kernel to be executed.  This region is not allowed to use values defined
 /// outside it.
 class LaunchOp : public Op<LaunchOp, OpTrait::AtLeastNOperands<6>::Impl,
-                           OpTrait::ZeroResult,
-                           OpTrait::NthRegionIsIsolatedAbove<0>::Impl> {
+                           OpTrait::ZeroResult, OpTrait::IsIsolatedFromAbove> {
 public:
   using Op::Op;
 

diff --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h
index 381a790..cc2eda8 100644 (file)
--- a/mlir/include/mlir/IR/Block.h
+++ b/mlir/include/mlir/IR/Block.h
@@ -415,7 +415,7 @@ public:
   /// Emit errors if `noteLoc` is provided; this location is used to point
   /// to the operation containing the region, the actual error is reported at
   /// the operation with an offending use.
-  bool isIsolatedAbove(llvm::Optional<Location> noteLoc = llvm::None);
+  bool isIsolatedFromAbove(llvm::Optional<Location> noteLoc = llvm::None);
 
   /// Walk the operations in this block in postorder, calling the callback for
   /// each operation.

diff --git a/mlir/include/mlir/IR/Function.h b/mlir/include/mlir/IR/Function.h
index 0ff93b6..933f6f8 100644 (file)
--- a/mlir/include/mlir/IR/Function.h
+++ b/mlir/include/mlir/IR/Function.h
@@ -328,7 +328,7 @@ private:
 /// allowed to implicitly capture global values, and all external references
 /// must use Function arguments or attributes.
 class FuncOp : public Op<FuncOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
-                         OpTrait::NthRegionIsIsolatedAbove<0>::Impl> {
+                         OpTrait::IsIsolatedFromAbove> {
 public:
   using Op::Op;
   static StringRef getOperationName() { return "func"; }

diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index ad62c59..c4257f0 100644 (file)
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -720,17 +720,36 @@ public:
   }
 };
 
-/// This verifiers that all operands used in N-th region of the given operation
-/// are defined within that region.
-template <unsigned RegionIdx> class NthRegionIsIsolatedAbove {
+/// This class provides the API for ops that are known to be isolated from
+/// above.
+template <typename ConcreteType>
+class IsIsolatedFromAbove
+    : public TraitBase<ConcreteType, IsIsolatedFromAbove> {
+public:
+  static AbstractOperation::OperationProperties getTraitProperties() {
+    return static_cast<AbstractOperation::OperationProperties>(
+        OperationProperty::IsolatedFromAbove);
+  }
+  static LogicalResult verifyTrait(Operation *op) {
+    for (auto &region : op->getRegions())
+      if (!region.isIsolatedFromAbove(op->getLoc()))
+        return failure();
+    return success();
+  }
+};
+
+/// This verifies that all operands used in N-th region of the given operation
+/// are defined within that region. This is a weaker variant to the
+/// 'IsIsolatedFromAbove' trait above as it applies only to one specific region.
+template <unsigned RegionIdx> class NthRegionIsIsolatedFromAbove {
 public:
   template <typename ConcreteType>
   class Impl : public TraitBase<ConcreteType,
-                                NthRegionIsIsolatedAbove<RegionIdx>::Impl> {
+                                NthRegionIsIsolatedFromAbove<RegionIdx>::Impl> {
   public:
     static LogicalResult verifyTrait(Operation *op) {
-      return op->getRegion(RegionIdx).isIsolatedAbove(op->getLoc()) ? success()
-                                                                    : failure();
+      return success(
+          op->getRegion(RegionIdx).isIsolatedFromAbove(op->getLoc()));
     }
   };
 };

diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 0a7a2aa..f974947 100644 (file)
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -386,6 +386,15 @@ public:
     return getTerminatorStatus() == TerminatorStatus::NonTerminator;
   }
 
+  /// Returns if the operation is known to be completely isolated from enclosing
+  /// regions, i.e. no internal regions reference values defined above this
+  /// operation.
+  bool isKnownIsolatedFromAbove() {
+    if (auto *absOp = getAbstractOperation())
+      return absOp->hasProperty(OperationProperty::IsolatedFromAbove);
+    return false;
+  }
+
   /// Attempt to fold this operation with the specified constant operand values
   /// - the elements in "operands" will correspond directly to the operands of
   /// the operation, but may be null if non-constant. If folding is successful,

diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index a5adad7..035421f 100644 (file)
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -74,6 +74,12 @@ enum class OperationProperty {
   /// This bit is set for an operation if it is a terminator: that means
   /// an operation at the end of a block.
   Terminator = 0x4,
+
+  /// This bit is set for operations that are completely isolated from above.
+  /// This is used for operations whose regions are explicit capture only, i.e.
+  /// they are never allowed to implicitly reference values defined above the
+  /// parent operation.
+  IsolatedFromAbove = 0x8,
 };
 
 /// This is a "type erased" representation of a registered operation.  This

diff --git a/mlir/lib/IR/Block.cpp b/mlir/lib/IR/Block.cpp
index 9595a72..0003fb2 100644 (file)
--- a/mlir/lib/IR/Block.cpp
+++ b/mlir/lib/IR/Block.cpp
@@ -413,7 +413,7 @@ static bool isRegionIsolatedAbove(Region &region, Region &limit,
   return true;
 }
 
-bool Region::isIsolatedAbove(llvm::Optional<Location> noteLoc) {
+bool Region::isIsolatedFromAbove(llvm::Optional<Location> noteLoc) {
   return isRegionIsolatedAbove(*this, *this, noteLoc);
 }