Introduce attributes that specify the final ABI for a spirv::ModuleOp.

To simplify the lowering into SPIR-V, while still respecting the ABI
requirements of SPIR-V/Vulkan, split the process into two
1) While lowering a function to SPIR-V (when the function is an entry
   point function), allow specifying attributes on arguments and
   function itself that describe the ABI of the function.
2) Add a pass that materializes the ABI described in the function.

Two attributes are needed.
1) Attribute on arguments of the entry point function that describe
   the descriptor_set, binding, storage class, etc, of the
   spv.globalVariable this argument will be replaced by
2) Attribute on function that specifies workgroup size, etc. (for now
   only workgroup size).

Add the pass -spirv-lower-abi-attrs to materialize the ABI described
by the attributes.

This change makes the SPIRVBasicTypeConverter class unnecessary and is
removed, further simplifying the SPIR-V lowering path.

PiperOrigin-RevId: 282387587

Author: Mahesh Ravishankar

include/mlir/Dialect/SPIRV/LayoutUtils.h

@@ -56,10 +56,10 @@ class VulkanLayoutUtils {
 public:
   using Size = uint64_t;
 
-  /// Returns a new type with layout info. Assigns the type size in bytes to the
-  /// `size`. Assigns the type alignment in bytes to the `alignment`.
-  static Type decorateType(spirv::StructType structType, Size &size,
-                           Size &alignment);
+  /// Returns a new StructType with layout info. Assigns the type size in bytes
+  /// to the `size`. Assigns the type alignment in bytes to the `alignment`.
+  static spirv::StructType decorateType(spirv::StructType structType,
+                                        Size &size, Size &alignment);
   /// Checks whether a type is legal in terms of Vulkan layout info
   /// decoration. A type is dynamically illegal if it's a composite type in the
   /// StorageBuffer, PhysicalStorageBuffer, Uniform, and PushConstant Storage

include/mlir/Dialect/SPIRV/Passes.h

@@ -27,14 +27,22 @@
 namespace mlir {
 namespace spirv {
 
-// Creates a module pass that converts composite types used by objects in the
-// StorageBuffer, PhysicalStorageBuffer, Uniform, and PushConstant storage
-// classes with layout information.
-//
-// Right now this pass only supports Vulkan layout rules.
+class ModuleOp;
+/// Creates a module pass that converts composite types used by objects in the
+/// StorageBuffer, PhysicalStorageBuffer, Uniform, and PushConstant storage
+/// classes with layout information.
+/// Right now this pass only supports Vulkan layout rules.
 std::unique_ptr<OpPassBase<mlir::ModuleOp>>
 createDecorateSPIRVCompositeTypeLayoutPass();
 
+/// Creates a module pass that lowers the ABI attributes specified during SPIR-V
+/// Lowering. Specifically,
+/// 1) Creates the global variables for arguments of entry point function using
+/// the specification in the ABI attributes for each argument.
+/// 2) Inserts the EntryPointOp and the ExecutionModeOp for entry point
+/// functions using the specification in the EntryPointAttr.
+std::unique_ptr<OpPassBase<spirv::ModuleOp>> createLowerABIAttributesPass();
+
 } // namespace spirv
 } // namespace mlir
 

include/mlir/Dialect/SPIRV/SPIRVLowering.h

@@ -23,35 +23,22 @@
 #define MLIR_DIALECT_SPIRV_SPIRVLOWERING_H
 
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/Support/StringExtras.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {
 
-/// Type conversion from Standard Types to SPIR-V Types.
-class SPIRVBasicTypeConverter : public TypeConverter {
-public:
-  /// Converts types to SPIR-V supported types.
-  virtual Type convertType(Type t);
-};
-
 /// Converts a function type according to the requirements of a SPIR-V entry
 /// function. The arguments need to be converted to spv.GlobalVariables of
 /// spv.ptr types so that they could be bound by the runtime.
 class SPIRVTypeConverter final : public TypeConverter {
 public:
-  explicit SPIRVTypeConverter(SPIRVBasicTypeConverter *basicTypeConverter)
-      : basicTypeConverter(basicTypeConverter) {}
+  using TypeConverter::TypeConverter;
 
   /// Converts types to SPIR-V types using the basic type converter.
   Type convertType(Type t) override;
-
-  /// Gets the basic type converter.
-  Type convertBasicType(Type t) { return basicTypeConverter->convertType(t); }
-
-private:
-  SPIRVBasicTypeConverter *basicTypeConverter;
 };
 
 /// Base class to define a conversion pattern to translate Ops into SPIR-V.
@@ -70,21 +57,35 @@ class SPIRVOpLowering : public OpConversionPattern<SourceOp> {
 private:
 };
 
+#include "mlir/Dialect/SPIRV/SPIRVLowering.h.inc"
+
 namespace spirv {
 /// Returns a value that represents a builtin variable value within the SPIR-V
 /// module.
 Value *getBuiltinVariableValue(Operation *op, spirv::BuiltIn builtin,
                                OpBuilder &builder);
 
 /// Legalizes a function as an entry function.
-LogicalResult lowerAsEntryFunction(FuncOp funcOp,
-                                   SPIRVTypeConverter *typeConverter,
-                                   ConversionPatternRewriter &rewriter,
-                                   FuncOp &newFuncOp);
-
-/// Finalizes entry function legalization. Inserts the spv.EntryPoint and
-/// spv.ExecutionMode ops.
-LogicalResult finalizeEntryFunction(FuncOp newFuncOp, OpBuilder &builder);
+FuncOp lowerAsEntryFunction(FuncOp funcOp, SPIRVTypeConverter &typeConverter,
+                            ConversionPatternRewriter &rewriter,
+                            ArrayRef<spirv::InterfaceVarABIAttr> argABIInfo,
+                            spirv::EntryPointABIAttr entryPointInfo);
+
+/// Attribute name for specifying argument ABI information.
+StringRef getInterfaceVarABIAttrName();
+
+/// Get the InterfaceVarABIAttr given its fields.
+InterfaceVarABIAttr getInterfaceVarABIAttr(unsigned descriptorSet,
+                                           unsigned binding,
+                                           spirv::StorageClass storageClass,
+                                           MLIRContext *context);
+
+/// Attribute name for specifying entry point information.
+StringRef getEntryPointABIAttrName();
+
+/// Get the EntryPointABIAttr given its fields.
+EntryPointABIAttr getEntryPointABIAttr(ArrayRef<int32_t> localSize,
+                                       MLIRContext *context);
 
 } // namespace spirv
 } // namespace mlir

include/mlir/Dialect/SPIRV/SPIRVLowering.td

@@ -0,0 +1,55 @@
+//===- SPIRVBase.td - MLIR SPIR-V Op Definitions Base file -*- tablegen -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This is the base file for supporting lowering to SPIR-V dialect. This
+// file defines SPIR-V attributes used for specifying the shader
+// interface or ABI. This is because SPIR-V module is expected to work in
+// an execution environment as specified by a client API. A SPIR-V module
+// needs to "link" correctly with the execution environment regarding the
+// resources that are used in the SPIR-V module and get populated with
+// data via the client API. The shader interface (or ABI) is passed into
+// SPIR-V lowering path via attributes defined in this file. A
+// compilation flow targeting SPIR-V is expected to attach such
+// attributes to resources and other suitable places.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SPIRV_LOWERING
+#define SPIRV_LOWERING
+
+include "mlir/Dialect/SPIRV/SPIRVBase.td"
+
+// For arguments that eventually map to spv.globalVariable for the
+// shader interface, this attribute specifies the information regarding
+// the global variable :
+// 1) Descriptor Set.
+// 2) Binding number.
+// 3) Storage class.
+def SPV_InterfaceVarABIAttr:
+    StructAttr<"InterfaceVarABIAttr", SPV_Dialect,
+               [StructFieldAttr<"descriptor_set", I32Attr>,
+                StructFieldAttr<"binding", I32Attr>,
+                StructFieldAttr<"storage_class", SPV_StorageClassAttr>]>;
+
+// For entry functions, this attribute specifies information related to entry
+// points in the generated SPIR-V module:
+// 1) WorkGroup Size.
+def SPV_EntryPointABIAttr:
+    StructAttr<"EntryPointABIAttr", SPV_Dialect,
+               [StructFieldAttr<"local_size", I32ElementsAttr>]>;
+
+#endif // SPIRV_LOWERING

include/mlir/Dialect/SPIRV/SPIRVStructureOps.td

@@ -243,7 +243,10 @@ def SPV_GlobalVariableOp : SPV_Op<"globalVariable", [InModuleScope, Symbol]> {
       "TypeAttr type, ArrayRef<NamedAttribute> namedAttrs", [{
       state.addAttribute("type", type);
       state.addAttributes(namedAttrs);
-    }]>
+    }]>,
+    OpBuilder<[{Builder *builder, OperationState &state,
+                Type type, StringRef name, unsigned descriptorSet,
+                unsigned binding}]>
   ];
 
   let results = (outs);

include/mlir/IR/OpBase.td

@@ -1230,6 +1230,10 @@ class ArrayMinCount<int n> : AttrConstraint<
     CPred<"$_self.cast<ArrayAttr>().size() >= " # n>,
     "with at least " # n # " elements">;
 
+class ArrayCount<int n> : AttrConstraint<
+    CPred<"$_self.cast<ArrayAttr>().size() == " #n>,
+    "with exactly " # n # " elements">;
+
 class IntArrayNthElemEq<int index, int value> : AttrConstraint<
     And<[
       CPred<"$_self.cast<ArrayAttr>().size() > " # index>,

lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

@@ -163,15 +163,26 @@ PatternMatchResult LaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
 PatternMatchResult
 KernelFnConversion::matchAndRewrite(FuncOp funcOp, ArrayRef<Value *> operands,
                                     ConversionPatternRewriter &rewriter) const {
-  FuncOp newFuncOp;
   if (!gpu::GPUDialect::isKernel(funcOp)) {
     return matchFailure();
   }
 
-  if (failed(spirv::lowerAsEntryFunction(funcOp, &typeConverter, rewriter,
-                                         newFuncOp))) {
+  SmallVector<spirv::InterfaceVarABIAttr, 4> argABI;
+  for (auto argNum : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
+    argABI.push_back(spirv::getInterfaceVarABIAttr(
+        0, argNum, spirv::StorageClass::StorageBuffer, rewriter.getContext()));
+  }
+  // TODO(ravishankarm) : For now set this to {32, 1, 1}. This is incorrect. The
+  // actual workgroup size needs to be plumbed through.
+  auto context = rewriter.getContext();
+  auto entryPointAttr = spirv::getEntryPointABIAttr({32, 1, 1}, context);
+  FuncOp newFuncOp = spirv::lowerAsEntryFunction(
+      funcOp, typeConverter, rewriter, argABI, entryPointAttr);
+  if (!newFuncOp) {
     return matchFailure();
   }
+  newFuncOp.removeAttr(Identifier::get(gpu::GPUDialect::getKernelFuncAttrName(),
+                                       rewriter.getContext()));
   return matchSuccess();
 }
 

lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.cpp

@@ -54,7 +54,7 @@ void GPUToSPIRVPass::runOnModule() {
     if (!gpu::GPUDialect::isKernel(funcOp)) {
       return;
     }
-    OpBuilder builder(module.getBodyRegion());
+    OpBuilder builder(funcOp.getOperation());
     // Create a new spirv::ModuleOp for this function, and clone the
     // function into it.
     // TODO : Generalize this to account for different extensions,
@@ -77,45 +77,20 @@ void GPUToSPIRVPass::runOnModule() {
   });
 
   /// Dialect conversion to lower the functions with the spirv::ModuleOps.
-  SPIRVBasicTypeConverter basicTypeConverter;
-  SPIRVTypeConverter typeConverter(&basicTypeConverter);
+  SPIRVTypeConverter typeConverter;
   OwningRewritePatternList patterns;
   populateGPUToSPIRVPatterns(context, typeConverter, patterns);
   populateStandardToSPIRVPatterns(context, typeConverter, patterns);
 
   ConversionTarget target(*context);
   target.addLegalDialect<spirv::SPIRVDialect>();
-  target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
-    // TODO(ravishankarm) : Currently lowering does not support handling
-    // function conversion of non-kernel functions. This is to be added.
-
-    // For kernel functions, verify that the signature is void(void).
-    return gpu::GPUDialect::isKernel(op) && op.getNumResults() == 0 &&
-           op.getNumArguments() == 0;
-  });
+  target.addDynamicallyLegalOp<FuncOp>(
+      [&](FuncOp op) { return typeConverter.isSignatureLegal(op.getType()); });
 
   if (failed(applyFullConversion(spirvModules, target, patterns,
                                  &typeConverter))) {
     return signalPassFailure();
   }
-
-  // After the SPIR-V modules have been generated, some finalization is needed
-  // for the entry functions. For example, adding spv.EntryPoint op,
-  // spv.ExecutionMode op, etc.
-  for (auto *spvModule : spirvModules) {
-    for (auto op :
-         cast<spirv::ModuleOp>(spvModule).getBlock().getOps<FuncOp>()) {
-      if (gpu::GPUDialect::isKernel(op)) {
-        OpBuilder builder(op.getContext());
-        builder.setInsertionPointAfter(op);
-        if (failed(spirv::finalizeEntryFunction(op, builder))) {
-          return signalPassFailure();
-        }
-        op.getOperation()->removeAttr(Identifier::get(
-            gpu::GPUDialect::getKernelFuncAttrName(), op.getContext()));
-      }
-    }
-  }
 }
 
 OpPassBase<ModuleOp> *createConvertGPUToSPIRVPass() {

lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

@@ -63,7 +63,7 @@ class ConstantIndexOpConversion final : public SPIRVOpLowering<ConstantOp> {
       return matchFailure();
     }
     auto spirvConstType =
-        typeConverter.convertBasicType(constIndexOp.getResult()->getType());
+        typeConverter.convertType(constIndexOp.getResult()->getType());
     auto spirvConstVal =
         rewriter.getIntegerAttr(spirvConstType, constAttr.getInt());
     rewriter.replaceOpWithNewOp<spirv::ConstantOp>(constIndexOp, spirvConstType,
@@ -120,7 +120,7 @@ class IntegerOpConversion final : public SPIRVOpLowering<StdOp> {
   matchAndRewrite(StdOp operation, ArrayRef<Value *> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto resultType =
-        this->typeConverter.convertBasicType(operation.getResult()->getType());
+        this->typeConverter.convertType(operation.getResult()->getType());
     rewriter.template replaceOpWithNewOp<SPIRVOp>(
         operation, resultType, operands, ArrayRef<NamedAttribute>());
     return this->matchSuccess();

lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.cpp

@@ -40,8 +40,7 @@ void ConvertStandardToSPIRVPass::runOnModule() {
   OwningRewritePatternList patterns;
   auto module = getModule();
 
-  SPIRVBasicTypeConverter basicTypeConverter;
-  SPIRVTypeConverter typeConverter(&basicTypeConverter);
+  SPIRVTypeConverter typeConverter;
   populateStandardToSPIRVPatterns(module.getContext(), typeConverter, patterns);
   ConversionTarget target(*(module.getContext()));
   target.addLegalDialect<spirv::SPIRVDialect>();