Rework module nesting in VM conversion. (#7873)

Introduces an explicit `vm.toplevel` attribute when setting up the nesting and then uses this as a signal that a nested module should be considered legal and left alone. The heuristic previously in use was to attempt to infer this based on parents. This turned out to be fragile if nesting a new module in an existing module (i.e. to extract a partial program and compile it without intending to create a layering issue).

I was going to clean the attribute up at the end but opted to leave it there as it was a good debugging aid as I tried to figure out what was going on.

Moved the one nesting test to a new nesting.mlir test case and reworked it to use a non-root based pipeline, which is more realistic for this scenario. Also, verified that what this test was doing was wrong (checks on "module" were inadvertently matching "vm.module", which was not the desired outcome). I believe the current tested behavior is desired.

iree/compiler/Dialect/HAL/Target/VMVX/test/smoketest.mlir

@@ -47,7 +47,7 @@ stream.executable public @add_dispatch_0 {
 //  CHECK-SAME:       interface = @io,
 //  CHECK-SAME:       ordinal = 0 : index
 //  CHECK-SAME:     }
-//       CHECK:     module {
+//       CHECK:     module attributes {vm.toplevel} {
 //  CHECK-NEXT:       vm.module public @module {
 //  CHECK-NEXT:         vm.func private @add_dispatch_0(
 //  CHECK-SAME:             %[[SCRATCHPAD:.+]]: !vm.buffer, %[[CONSTANTS:.+]]: !vm.buffer,

iree/compiler/Dialect/VM/Conversion/ConversionTarget.cpp

@@ -24,18 +24,25 @@ VMConversionTarget::nestModuleForConversion(mlir::ModuleOp outerModuleOp) {
     outerModuleOp.getBodyRegion().getBlocks().push_back(new Block());
     outerModuleOp.push_back(innerModuleOp);
   }
+
+  outerModuleOp->setAttr("vm.toplevel",
+                         UnitAttr::get(outerModuleOp.getContext()));
   return std::make_pair(outerModuleOp, innerModuleOp);
 }
 
+// static
+bool VMConversionTarget::isTopLevelModule(mlir::ModuleOp moduleOp) {
+  return !moduleOp->getParentOp() || moduleOp->hasAttr("vm.toplevel");
+}
+
 VMConversionTarget::VMConversionTarget(MLIRContext *context)
     : ConversionTarget(*context) {
   addLegalDialect<IREE::VM::VMDialect>();
 
   // NOTE: we need to allow the outermost std.module to be legal to support the
   // double-nesting (module { vm.module { ... } }).
-  addDynamicallyLegalOp<mlir::ModuleOp>(+[](mlir::ModuleOp op) {
-    return !op->getParentOp() || !isa<ModuleOp>(op->getParentOp());
-  });
+  addDynamicallyLegalOp<mlir::ModuleOp>(
+      +[](mlir::ModuleOp op) { return isTopLevelModule(op); });
 }
 
 }  // namespace iree_compiler

iree/compiler/Dialect/VM/Conversion/ConversionTarget.h

@@ -27,10 +27,14 @@ class VMConversionTarget : public ConversionTarget {
   // Example:
   //  module { func @foo() { ... } }
   // ->
-  //  module { module { func @foo() { ... } } }
+  //  module attributes {vm.toplevel} { module { func @foo() { ... } } }
   static std::pair<mlir::ModuleOp, mlir::ModuleOp> nestModuleForConversion(
       mlir::ModuleOp outerModuleOp);
 
+  // Returns whether this is the outer module as setup via
+  // nestModuleForConversion. Use for patterns which need to distinguish.
+  static bool isTopLevelModule(mlir::ModuleOp moduleOp);
+
   VMConversionTarget(MLIRContext *context);
 };
 

iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVM.cpp

@@ -7,6 +7,7 @@
 #include "iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVM.h"
 
 #include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
+#include "iree/compiler/Dialect/VM/Conversion/ConversionTarget.h"
 #include "iree/compiler/Dialect/VM/Conversion/TargetOptions.h"
 #include "iree/compiler/Dialect/VM/Conversion/TypeConverter.h"
 #include "iree/compiler/Dialect/VM/IR/VMOps.h"
@@ -32,7 +33,7 @@ class ModuleOpConversion : public OpConversionPattern<ModuleOp> {
       ConversionPatternRewriter &rewriter) const override {
     // Do not attempt to convert the top level module.
     // This mechanism can only support rewriting non top-level modules.
-    if (!srcOp->getParentOp() || !isa<ModuleOp>(srcOp->getParentOp())) {
+    if (VMConversionTarget::isTopLevelModule(srcOp)) {
       return failure();
     }
 

iree/compiler/Dialect/VM/Conversion/StandardToVM/test/BUILD

@@ -24,6 +24,7 @@ iree_lit_test_suite(
             "control_flow_ops.mlir",
             "conversion_ops.mlir",
             "func_attrs.mlir",
+            "nesting.mlir",
             "structural_ops.mlir",
         ],
         include = ["*.mlir"],

iree/compiler/Dialect/VM/Conversion/StandardToVM/test/CMakeLists.txt

@@ -21,6 +21,7 @@ iree_lit_test_suite(
     "control_flow_ops.mlir"
     "conversion_ops.mlir"
     "func_attrs.mlir"
+    "nesting.mlir"
     "structural_ops.mlir"
   DATA
     iree::tools::IreeFileCheck

iree/compiler/Dialect/VM/Conversion/StandardToVM/test/nesting.mlir

@@ -0,0 +1,21 @@
+// RUN: iree-opt -split-input-file -pass-pipeline='builtin.module(test-iree-convert-std-to-vm)' %s | IreeFileCheck %s
+
+// Note that checks are ambiguous between "module" and "vm.module" so we rely
+// on vm.module printing as `vm.module public @foo`
+
+// CHECK-LABEL: module @outerBuiltinModule
+module @outerBuiltinModule {
+  // CHECK-NEXT: module @innerBuiltinModule attributes {vm.toplevel}
+  module @innerBuiltinModule attributes {vm.toplevel} {
+    // CHECK-NEXT: vm.module public @outerVmModule
+    module @outerVmModule {
+      // CHECK-NEXT: vm.module public @deeplyNested
+      module @deeplyNested {
+        // CHECK: vm.func private @foo
+        func @foo() {
+          return
+        }
+      }
+    }
+  }
+}

iree/compiler/Dialect/VM/Conversion/StandardToVM/test/structural_ops.mlir

@@ -40,22 +40,3 @@ module {
 }
 
 }
-
-// -----
-
-// CHECK: module
-module {
-  // CHECK: module
-  module {
-    // CHECK: module
-    module {
-      // CHECK-LABEL: vm.module public @deeplyNested
-      module @deeplyNested {
-        // CHECK: vm.func private @foo
-        func @foo() {
-          return
-        }
-      }
-    }
-  }
-}