Remove symbolTable op arg and simplify tests.

Author: ingomueller-net

mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td

@@ -153,7 +153,7 @@ def MemRefAllocaToGlobalOp :
       DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let description = [{
     Inserts a new `memref.global` for each provided `memref.alloca` into the
-    provided symbol table (e.g., a `builtin.module`) and replaces it with a
+    nearest symbol table (e.g., a `builtin.module`) and replaces it with a
     `memref.get_global`. This is useful, for example, for allocations that
     should reside in the shared memory of a GPU, which have to be declared as
     globals.
@@ -164,8 +164,8 @@ def MemRefAllocaToGlobalOp :
 
     ```mlir
     %get_global, %global =
-        transform.memref.alloca_to_global %alloca in %module
-          : (!transform.op<"builtin.module">, !transform.op<"memref.alloca">)
+        transform.memref.alloca_to_global %alloca
+          : (!transform.op<"memref.alloca">)
             -> (!transform.any_op, !transform.any_op)
     ```
 
@@ -195,20 +195,16 @@ def MemRefAllocaToGlobalOp :
 
     #### Return modes
 
-    Emits a definite failure if not exactly one symbol table payload op was
-    provided or any of the `alloca` payload ops is not inside that symbol table
-    op, and succeeds otherwise. The returned handles refer to the
-    `memref.get_global` and `memref.global` ops that were inserted by the
-    transformation.
+    Succeeds always. The returned handles refer to the `memref.get_global` and
+    `memref.global` ops that were inserted by the transformation.
   }];
 
-  let arguments = (ins TransformHandleTypeInterface:$symbolTable,
-                   Transform_MemRefAllocaOp:$alloca);
+  let arguments = (ins Transform_MemRefAllocaOp:$alloca);
   let results = (outs TransformHandleTypeInterface:$getGlobal,
                   TransformHandleTypeInterface:$global);
 
   let assemblyFormat = [{
-    $alloca `in` $symbolTable attr-dict `:` functional-type(operands, results)
+    $alloca attr-dict `:` functional-type(operands, results)
   }];
 }
 

mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp

@@ -139,26 +139,6 @@ transform::MemRefAllocaToGlobalOp::apply(transform::TransformRewriter &rewriter,
   SmallVector<memref::GlobalOp> globalOps;
   SmallVector<memref::GetGlobalOp> getGlobalOps;
 
-  // Get containing symbol table op.
-  auto symbolTableOps = state.getPayloadOps(getSymbolTable());
-  if (!llvm::hasSingleElement(symbolTableOps)) {
-    return emitDefiniteFailure()
-           << Twine("expected exactly one 'symbolTable' payload, but found ") +
-                  std::to_string(llvm::range_size(symbolTableOps));
-  }
-  Operation *symbolTableOp = *symbolTableOps.begin();
-  if (!symbolTableOp->hasTrait<OpTrait::SymbolTable>()) {
-    return emitDefiniteFailure() << Twine(
-               "expected 'symbolTable' payload to have 'SymbolTable' trait");
-  }
-  SymbolTable symbolTable(symbolTableOp);
-
-  {
-    size_t numAllocaOps = llvm::range_size(allocaOps);
-    globalOps.reserve(numAllocaOps);
-    getGlobalOps.reserve(numAllocaOps);
-  }
-
   // Transform `memref.alloca`s.
   for (auto *op : allocaOps) {
     auto alloca = cast<memref::AllocaOp>(op);
@@ -167,14 +147,15 @@ transform::MemRefAllocaToGlobalOp::apply(transform::TransformRewriter &rewriter,
 
     memref::GlobalOp globalOp;
     {
+      // Find nearest symbol table.
+      Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(op);
+      assert(symbolTableOp && "expected alloca payload to be in symbol table");
+      SymbolTable symbolTable(symbolTableOp);
+
       // Insert a `memref.global` into the symbol table.
-      if (symbolTable.getOp() != SymbolTable::getNearestSymbolTable(op)) {
-        return emitDefiniteFailure() << "expected 'alloca' payload to be "
-                                        "inside 'symbolTable' payload";
-      }
       Type resultType = alloca.getResult().getType();
-      // TODO: Add a better builder for this.
       OpBuilder builder(rewriter.getContext());
+      // TODO: Add a better builder for this.
       globalOp = builder.create<memref::GlobalOp>(
           loc, StringAttr::get(ctx, "alloca"), StringAttr::get(ctx, "private"),
           TypeAttr::get(resultType), Attribute{}, UnitAttr{}, IntegerAttr{});
@@ -200,7 +181,6 @@ transform::MemRefAllocaToGlobalOp::apply(transform::TransformRewriter &rewriter,
 
 void transform::MemRefAllocaToGlobalOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  onlyReadsHandle(getSymbolTable(), effects);
   producesHandle(getGlobal(), effects);
   producesHandle(getGetGlobal(), effects);
   consumesHandle(getAlloca(), effects);

mlir/test/Dialect/MemRef/transform-ops.mlir

@@ -23,16 +23,9 @@ func.func @func(%lb: index, %ub: index) {
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !transform.any_op):
   %alloca = transform.structured.match ops{["memref.alloca"]} in %arg0
-      : (!transform.any_op) -> !transform.any_op
-  %module = transform.structured.match ops{["builtin.module"]} in %arg0
-      : (!transform.any_op) -> !transform.any_op
-  %alloca_typed = transform.cast %alloca
-      : !transform.any_op to !transform.op<"memref.alloca">
-  %module_typed = transform.cast %module
-      : !transform.any_op to !transform.op<"builtin.module">
-  %get_global, %global =
-      transform.memref.alloca_to_global %alloca_typed in %module_typed
-        : (!transform.op<"builtin.module">, !transform.op<"memref.alloca">)
+      : (!transform.any_op) -> !transform.op<"memref.alloca">
+  %get_global, %global = transform.memref.alloca_to_global %alloca
+        : (!transform.op<"memref.alloca">)
           -> (!transform.any_op, !transform.any_op)
 }