[mlir][spirv] Relax verification to allow flexible placement

[antiagainst]

Summary

Thus far certain SPIR-V ops have been required to be in spv.module.
While this provides strong verification to catch unexpected errors,
it's quite rigid and makes progressive lowering difficult. Sometimes
we would like to partially lower ops from other dialects, which may
involve creating ops like global variables that should be placed in
other module-like ops. So this commit relaxes the requirement of
such SPIR-V ops' scope to module-like ops. Similarly for function-
like ops.

[merge_guards_bot]

Unit tests: pass. 62195 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[antiagainst]

Change more ops' verification

[merge_guards_bot]

Unit tests: pass. 62195 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td

 
 // TODO(ravishankarm): From 1.4, this should also include Composite type.
 def SPV_SelectType : AnyTypeOf<[SPV_Scalar, SPV_Vector, SPV_AnyPtr]>;
 
 //===----------------------------------------------------------------------===//
 // SPIR-V OpTrait definitions
 //===----------------------------------------------------------------------===//
 
-// Check that an op can only be used within the scope of a FuncOp.
+// Check that an op can only be used within the scope of a function-like op.
 def InFunctionScope : PredOpTrait<
-  "op must appear in a 'func' block",
-  CPred<"($_op.getParentOfType<FuncOp>())">>;
+  "op must appear in a function-like op's block",
+  CPred<"isNestedInFunctionLikeOp($_op.getParentOp())">>;
 
-// Check that an op can only be used within the scope of a SPIR-V ModuleOp.
+// Check that an op can only be used within the scope of a module-like op.
 def InModuleScope : PredOpTrait<
-  "op must appear in a 'spv.module' block",
-  CPred<"llvm::isa_and_nonnull<spirv::ModuleOp>($_op.getParentOp())">>;
+  "op must appear in a module-like op's block",
+  CPred<"isDirectInModuleLikeOp($_op.getParentOp())">>;
 
 //===----------------------------------------------------------------------===//
 // SPIR-V opcode specification
 //===----------------------------------------------------------------------===//
 
 class SPV_OpCode<string name, int val> {
   // Name used as reference to retrieve the opcode
   string opname = name;

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

 static constexpr const char kValueAttrName[] = "value";
 static constexpr const char kValuesAttrName[] = "values";
 static constexpr const char kVariableAttrName[] = "variable";
 
 //===----------------------------------------------------------------------===//
 // Common utility functions
 //===----------------------------------------------------------------------===//
 
+/// Returns true if the given op is a function-like op or nested in a
+/// function-like op without a module-like op in the middle.
+static bool isNestedInFunctionLikeOp(Operation *op) {
+  if (!op)
+    return false;
+  if (op->hasTrait<OpTrait::SymbolTable>())
+    return false;
+  if (op->hasTrait<OpTrait::FunctionLike>())
+    return true;
+  return isNestedInFunctionLikeOp(op->getParentOp());
+}
+
+/// Returns true if the given op is an module-like op that maintains a symbol
+/// table.
+static bool isDirectInModuleLikeOp(Operation *op) {
+  return op && op->hasTrait<OpTrait::SymbolTable>();
+}
+
 static LogicalResult extractValueFromConstOp(Operation *op, int32_t &value) {
   auto constOp = dyn_cast_or_null<spirv::ConstantOp>(op);
   if (!constOp) {
     return failure();
   }
   auto valueAttr = constOp.value();
   auto integerValueAttr = valueAttr.dyn_cast<IntegerAttr>();
   if (!integerValueAttr) {
static void print(spirv::AddressOfOp addressOfOp, OpAsmPrinter &printer) {
   printer << ' ';
   printer.printSymbolName(addressOfOp.variable());
 
   // Print the type.
   printer << " : " << addressOfOp.pointer().getType();
 }
 
 static LogicalResult verify(spirv::AddressOfOp addressOfOp) {
-  auto moduleOp = addressOfOp.getParentOfType<spirv::ModuleOp>();
-  auto varOp =
-      moduleOp.lookupSymbol<spirv::GlobalVariableOp>(addressOfOp.variable());
+  auto varOp = dyn_cast_or_null<spirv::GlobalVariableOp>(
+      SymbolTable::lookupNearestSymbolFrom(addressOfOp.getParentOp(),
+                                           addressOfOp.variable()));
   if (!varOp) {
     return addressOfOp.emitOpError("expected spv.globalVariable symbol");
   }
   if (addressOfOp.pointer().getType() != varOp.type()) {
     return addressOfOp.emitOpError(
         "result type mismatch with the referenced global variable's type");
   }
   return success();
static void print(spirv::FunctionCallOp functionCallOp, OpAsmPrinter &printer) {
   printer << spirv::FunctionCallOp::getOperationName() << ' '
           << functionCallOp.getAttr(kCallee) << '('
           << functionCallOp.arguments() << ") : " << functionType;
 }
 
 static LogicalResult verify(spirv::FunctionCallOp functionCallOp) {
   auto fnName = functionCallOp.callee();
 
-  auto moduleOp = functionCallOp.getParentOfType<spirv::ModuleOp>();
-  if (!moduleOp) {
-    return functionCallOp.emitOpError(
-        "must appear in a function inside 'spv.module'");
-  }
-
-  auto funcOp = moduleOp.lookupSymbol<FuncOp>(fnName);
+  auto funcOp = dyn_cast_or_null<FuncOp>(SymbolTable::lookupNearestSymbolFrom(
+      functionCallOp.getParentOp(), fnName));
   if (!funcOp) {
     return functionCallOp.emitOpError("callee function '")
-           << fnName << "' not found in 'spv.module'";
+           << fnName << "' not found in nearest symbol table";
   }
 
   auto functionType = funcOp.getType();
 
   if (functionCallOp.getNumResults() > 1) {
     return functionCallOp.emitOpError(
                "expected callee function to have 0 or 1 result, but provided ")
            << functionCallOp.getNumResults();
static LogicalResult verify(spirv::GlobalVariableOp varOp) {
   // SPIR-V spec: "Storage Class is the Storage Class of the memory holding the
   // object. It cannot be Generic. It must be the same as the Storage Class
   // operand of the Result Type."
   if (varOp.storageClass() == spirv::StorageClass::Generic)
     return varOp.emitOpError("storage class cannot be 'Generic'");
 
   if (auto init =
           varOp.getAttrOfType<FlatSymbolRefAttr>(kInitializerAttrName)) {
-    auto moduleOp = varOp.getParentOfType<spirv::ModuleOp>();
-    auto *initOp = moduleOp.lookupSymbol(init.getValue());
+    Operation *initOp = SymbolTable::lookupNearestSymbolFrom(
+        varOp.getParentOp(), init.getValue());
     // TODO: Currently only variable initialization with specialization
     // constants and other variables is supported. They could be normal
     // constants in the module scope as well.
     if (!initOp || !(isa<spirv::GlobalVariableOp>(initOp) ||
                      isa<spirv::SpecConstantOp>(initOp))) {
       return varOp.emitOpError("initializer must be result of a "
                                "spv.specConstant or spv.globalVariable op");
     }

 
 static void print(spirv::ReferenceOfOp referenceOfOp, OpAsmPrinter &printer) {
   printer << spirv::ReferenceOfOp::getOperationName() << ' ';
   printer.printSymbolName(referenceOfOp.spec_const());
   printer << " : " << referenceOfOp.reference().getType();
 }
 
 static LogicalResult verify(spirv::ReferenceOfOp referenceOfOp) {
-  auto moduleOp = referenceOfOp.getParentOfType<spirv::ModuleOp>();
-  auto specConstOp =
-      moduleOp.lookupSymbol<spirv::SpecConstantOp>(referenceOfOp.spec_const());
+  auto specConstOp = dyn_cast_or_null<spirv::SpecConstantOp>(
+      SymbolTable::lookupNearestSymbolFrom(referenceOfOp.getParentOp(),
+                                           referenceOfOp.spec_const()));
   if (!specConstOp) {
     return referenceOfOp.emitOpError("expected spv.specConstant symbol");
   }
   if (referenceOfOp.reference().getType() !=
       specConstOp.default_value().getType()) {
     return referenceOfOp.emitOpError("result type mismatch with the referenced "
                                      "specialization constant's type");
   }

mlir/test/Dialect/SPIRV/control-flow-ops.mlir

spv.module "Logical" "GLSL450" {
 
   func @f_3(%arg0 : i32) -> (i32) {
     spv.ReturnValue %arg0 : i32
   }
 }
 
 // -----
 
+// Allow calling functions in other module-like ops
+func @callee() {
+  spv.Return
+}
+
+func @caller() {
+  // CHECK: spv.FunctionCall
+  spv.FunctionCall @callee() : () -> ()
+  spv.Return
+}
+
+// -----
+
 spv.module "Logical" "GLSL450" {
   func @f_invalid_result_type(%arg0 : i32, %arg1 : i32) -> () {
     // expected-error @+1 {{expected callee function to have 0 or 1 result, but provided 2}}
     %0 = spv.FunctionCall @f_invalid_result_type(%arg0, %arg1) : (i32, i32) -> (i32, i32)
     spv.Return
   }
 }
 
func @f_type_mismatch(%arg0 : i32, %arg1 : i32) -> i32 {
     spv.Return
   }
 }
 
 // -----
 
 spv.module "Logical" "GLSL450" {
   func @f_foo(%arg0 : i32, %arg1 : i32) -> i32 {
-    // expected-error @+1 {{op callee function 'f_undefined' not found in 'spv.module'}}
+    // expected-error @+1 {{op callee function 'f_undefined' not found in nearest symbol table}}
     %0 = spv.FunctionCall @f_undefined(%arg0, %arg0) : (i32, i32) -> i32
     spv.Return
   }
 }
 
 // -----
 
-func @f_foo(%arg0 : i32, %arg1 : i32) -> i32 {
-    // expected-error @+1 {{must appear in a function inside 'spv.module'}}
-    %0 = spv.FunctionCall @f_foo(%arg0, %arg0) : (i32, i32) -> i32
-    spv.Return
-}
-
-// -----
-
 //===----------------------------------------------------------------------===//
 // spv.loop
 //===----------------------------------------------------------------------===//
 
 // for (int i = 0; i < count; ++i) {}
 func @loop(%count : i32) -> () {
   %zero = spv.constant 0: i32
   %one = spv.constant 1: i32
func @in_loop(%cond : i1) -> () {
   ^merge:
     spv._merge
   }
   spv.Return
 }
 
 // -----
 
+// CHECK-LABEL: in_other_func_like_op
+func @in_other_func_like_op() {
+  // CHECK: spv.Return 
+  spv.Return
+}
+
+// -----
+
 "foo.function"() ({
-  // expected-error @+1 {{op must appear in a 'func' block}}
+  // expected-error @+1 {{op must appear in a function-like op's block}}
   spv.Return
 })  : () -> ()
 
 // -----
 
 // Return mismatches function signature
 spv.module "Logical" "GLSL450" {
   func @work() -> (i32) {
func @in_loop(%cond : i1) -> (i32) {
   %one = spv.constant 1 : i32
   spv.ReturnValue %one : i32
 }
 
 // -----
 
 "foo.function"() ({
   %0 = spv.constant true
-  // expected-error @+1 {{op must appear in a 'func' block}}
+  // expected-error @+1 {{op must appear in a function-like op's block}}
   spv.ReturnValue %0 : i1
 })  : () -> ()
 
 // -----
 
 func @value_count_mismatch() -> () {
   %0 = spv.constant 42 : i32
   // expected-error @+1 {{op returns 1 value but enclosing function requires 0 results}}

mlir/test/Dialect/SPIRV/structure-ops.mlir

func @access_chain() -> () {
     %1 = spv._address_of @var1 : !spv.ptr<!spv.struct<f32, !spv.array<4xf32>>, Input>
     %2 = spv.AccessChain %1[%0, %0] : !spv.ptr<!spv.struct<f32, !spv.array<4xf32>>, Input>
     spv.Return
   }
 }
 
 // -----
 
+// Allow taking address of global variables in other module-like ops
+spv.globalVariable @var : !spv.ptr<!spv.struct<f32, !spv.array<4xf32>>, Input>
+func @address_of() -> () {
+  // CHECK: spv._address_of @var
+  %1 = spv._address_of @var : !spv.ptr<!spv.struct<f32, !spv.array<4xf32>>, Input>
+  return
+}
+
+// -----
+
 spv.module "Logical" "GLSL450" {
   spv.globalVariable @var1 : !spv.ptr<!spv.struct<f32, !spv.array<4xf32>>, Input>
   func @foo() -> () {
     // expected-error @+1 {{expected spv.globalVariable symbol}}
     %0 = spv._address_of @var2 : !spv.ptr<!spv.struct<f32, !spv.array<4xf32>>, Input>
   }
 }
 

 /// when interface entries of EntryPointOp are not
 /// spv.Variables. There is currently no other op that has a spv.ptr
 /// return type
 
 // -----
 
 spv.module "Logical" "GLSL450" {
    func @do_nothing() -> () {
-     // expected-error @+1 {{'spv.EntryPoint' op failed to verify that op must appear in a 'spv.module' block}}
+     // expected-error @+1 {{op must appear in a module-like op's block}}
      spv.EntryPoint "GLCompute" @do_something
    }
 }
 
 // -----
 
 spv.module "Logical" "GLSL450" {
    func @do_nothing() -> () {

 spv.module "Logical" "GLSL450" {
   spv.specConstant @sc = 4.0 : f32
   // CHECK: spv.globalVariable @var initializer(@sc) : !spv.ptr<f32, Private>
   spv.globalVariable @var initializer(@sc) : !spv.ptr<f32, Private>
 }
 
 // -----
 
+// Allow initializers coming from other module-like ops
+spv.specConstant @sc = 4.0 : f32
+// CHECK: spv.globalVariable @var initializer(@sc)
+spv.globalVariable @var initializer(@sc) : !spv.ptr<f32, Private>
+
+// -----
+
 spv.module "Logical" "GLSL450" {
   // CHECK: spv.globalVariable @var0 bind(1, 2) : !spv.ptr<f32, Uniform>
   spv.globalVariable @var0 bind(1, 2) : !spv.ptr<f32, Uniform>
 }
 
 // TODO: Fix test case after initialization with constant is addressed
 // spv.module "Logical" "GLSL450" {
 //   %0 = spv.constant 4.0 : f32
 //   // CHECK1: spv.globalVariable @var1 initializer(%0) {binding = 5 : i32} : !spv.ptr<f32, Private>
 //   spv.globalVariable @var1 initializer(%0) {binding = 5 : i32} : !spv.ptr<f32, Private>
 // }
 
 // -----
 
 spv.module "Logical" "GLSL450" {
   // CHECK: spv.globalVariable @var1 built_in("GlobalInvocationID") : !spv.ptr<vector<3xi32>, Input>
   spv.globalVariable @var1 built_in("GlobalInvocationID") : !spv.ptr<vector<3xi32>, Input>
   // CHECK: spv.globalVariable @var2 built_in("GlobalInvocationID") : !spv.ptr<vector<3xi32>, Input>
   spv.globalVariable @var2 {built_in = "GlobalInvocationID"} : !spv.ptr<vector<3xi32>, Input>
 }
 
 // -----
 
+// Allow in other module-like ops
+module {
+  // CHECK: spv.globalVariable
+  spv.globalVariable @var0 : !spv.ptr<f32, Input>
+}
+
+// -----
+
 spv.module "Logical" "GLSL450" {
   // expected-error @+1 {{expected spv.ptr type}}
   spv.globalVariable @var0 : f32
 }
 
 // -----
 
 spv.module "Logical" "GLSL450" {
   // expected-error @+1 {{op initializer must be result of a spv.specConstant or spv.globalVariable op}}
   spv.globalVariable @var0 initializer(@var1) : !spv.ptr<f32, Private>
 }
 
 // -----
 
 spv.module "Logical" "GLSL450" {
   // expected-error @+1 {{storage class cannot be 'Generic'}}
   spv.globalVariable @var0 : !spv.ptr<f32, Generic>
 }
 
 // -----
 
 spv.module "Logical" "GLSL450" {
   func @foo() {
-    // expected-error @+1 {{op failed to verify that op must appear in a 'spv.module' block}}
+    // expected-error @+1 {{op must appear in a module-like op's block}}
     spv.globalVariable @var0 : !spv.ptr<f32, Input>
     spv.Return
   }
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//

 
 // -----
 
 //===----------------------------------------------------------------------===//
 // spv._module_end
 //===----------------------------------------------------------------------===//
 
 func @module_end_not_in_module() -> () {
-  // expected-error @+1 {{op must appear in a 'spv.module' block}}
+  // expected-error @+1 {{op must appear in a module-like op's block}}
   spv._module_end
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // spv._reference_of
 //===----------------------------------------------------------------------===//
func @compute() -> f32 {
     %1 = spv.constant 6.0 : f32
     %2 = spv.FAdd %0, %1 : f32
     spv.ReturnValue %2 : f32
   }
 }
 
 // -----
 
+// Allow taking reference of spec constant in other module-like ops
+spv.specConstant @sc = 5 : i32
+func @reference_of() {
+  // CHECK: spv._reference_of @sc
+  %0 = spv._reference_of @sc : i32
+  return
+}
+
+// -----
+
 spv.module "Logical" "GLSL450" {
   func @foo() -> () {
     // expected-error @+1 {{expected spv.specConstant symbol}}
     %0 = spv._reference_of @sc : i32
     spv.Return
   }
 }
 

 spv.module "Logical" "GLSL450" {
   // expected-error @+1 {{default value can only be a bool, integer, or float scalar}}
   spv.specConstant @sc = dense<[2, 3]> : vector<2xi32>
 }
 
 // -----
 
 func @use_in_function() -> () {
-  // expected-error @+1 {{op must appear in a 'spv.module' block}}
+  // expected-error @+1 {{op must appear in a module-like op's block}}
   spv.specConstant @sc = false
   return
 }