Commit 399887c9 authored by Lei Zhang's avatar Lei Zhang
Browse files

[mlir][spirv] Add resource limits into target environment 

This commit adds two resource limits, max_compute_workgroup_size
and max_compute_workgroup_invocations as resource limits to
the target environment. They are not used at the current moment,
but they will affect the SPIR-V CodeGen. Adding for now to have
a proper target environment modelling.

Differential Revision: https://reviews.llvm.org/D73905
parent 2926a651

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h

+4 −0
Original line number Diff line number Diff line
@@ -43,6 +43,10 @@ StringRef getEntryPointABIAttrName(); 
EntryPointABIAttr getEntryPointABIAttr(ArrayRef<int32_t> localSize,

                                       MLIRContext *context);



/// Returns a default resource limits attribute that uses numbers from

/// "Table 46. Required Limits" of the Vulkan spec.

ResourceLimitsAttr getDefaultResourceLimits(MLIRContext *context);



/// Returns the attribute name for specifying SPIR-V target environment.

StringRef getTargetEnvAttrName();

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td

+14 −2
Original line number Diff line number Diff line
@@ -48,12 +48,24 @@ def SPV_ExtensionArrayAttr : TypedArrayAttrBase< 
def SPV_CapabilityArrayAttr : TypedArrayAttrBase<

    SPV_CapabilityAttr, "SPIR-V capability array attribute">;



// This attribute specifies the limits for various resources on the target

// architecture.

//

// See https://renderdoc.org/vkspec_chunked/chap36.html#limits for the complete

// list of limits and their explanation for the Vulkan API. The following ones

// are those affecting SPIR-V CodeGen.

def SPV_ResourceLimitsAttr : StructAttr<"ResourceLimitsAttr", SPV_Dialect, [

    StructFieldAttr<"max_compute_workgroup_invocations", I32Attr>,

    StructFieldAttr<"max_compute_workgroup_size", I32ElementsAttr>

]>;



// For the generated SPIR-V module, this attribute specifies the target version,

// allowed extensions and capabilities.

// allowed extensions and capabilities, and resource limits.

def SPV_TargetEnvAttr : StructAttr<"TargetEnvAttr", SPV_Dialect, [

    StructFieldAttr<"version", SPV_VersionAttr>,

    StructFieldAttr<"extensions", SPV_ExtensionArrayAttr>,

    StructFieldAttr<"capabilities", SPV_CapabilityArrayAttr>

    StructFieldAttr<"capabilities", SPV_CapabilityArrayAttr>,

    StructFieldAttr<"limits", SPV_ResourceLimitsAttr>

]>;



#endif // SPIRV_TARGET_AND_ABI

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

+2 −2
Original line number Diff line number Diff line
@@ -664,8 +664,8 @@ LogicalResult SPIRVDialect::verifyOperationAttribute(Operation *op, 
             << symbol

             << "' must be a dictionary attribute containing one 32-bit "

                "integer attribute 'version', one string array attribute "

                "'extensions', and one 32-bit integer array attribute "

                "'capabilities'";

                "'extensions', one 32-bit integer array attribute "

                "'capabilities', and one dictionary attribute 'limits'";

  } else {

    return op->emitError("found unsupported '")

           << symbol << "' attribute on operation";

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

+12 −1
Original line number Diff line number Diff line
@@ -45,6 +45,17 @@ spirv::getEntryPointABIAttr(ArrayRef<int32_t> localSize, MLIRContext *context) { 
      context);

}



spirv::ResourceLimitsAttr

spirv::getDefaultResourceLimits(MLIRContext *context) {

  auto i32Type = IntegerType::get(32, context);

  auto v3i32Type = VectorType::get(3, i32Type);



  // These numbers are from "Table 46. Required Limits" of the Vulkan spec.

  return spirv::ResourceLimitsAttr ::get(

      IntegerAttr::get(i32Type, 128),

      DenseIntElementsAttr::get<int32_t>(v3i32Type, {128, 128, 64}), context);

}



StringRef spirv::getTargetEnvAttrName() { return "spv.target_env"; }



spirv::TargetEnvAttr spirv::getDefaultTargetEnv(MLIRContext *context) {
@@ -54,7 +65,7 @@ spirv::TargetEnvAttr spirv::getDefaultTargetEnv(MLIRContext *context) { 
      builder.getI32ArrayAttr({}),

      builder.getI32ArrayAttr(

          {static_cast<uint32_t>(spirv::Capability::Shader)}),

      context);

      spirv::getDefaultResourceLimits(context), context);

}



spirv::TargetEnvAttr spirv::lookupTargetEnvOrDefault(Operation *op) {

mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir

+20 −4
Original line number Diff line number Diff line
@@ -18,7 +18,11 @@ module attributes { 
  spv.target_env = {

    version = 3 : i32,

    extensions = [],

    capabilities = [1: i32, 63: i32] // Shader, GroupNonUniformArithmetic

    capabilities = [1: i32, 63: i32], // Shader, GroupNonUniformArithmetic

    limits = {

      max_compute_workgroup_invocations = 128 : i32,

      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>

    }

  }

} {
@@ -79,7 +83,11 @@ module attributes { 
  spv.target_env = {

    version = 3 : i32,

    extensions = [],

    capabilities = [1: i32, 63: i32] // Shader, GroupNonUniformArithmetic

    capabilities = [1: i32, 63: i32], // Shader, GroupNonUniformArithmetic

    limits = {

      max_compute_workgroup_invocations = 128 : i32,

      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>

    }

  }

} {

func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>) {
@@ -111,7 +119,11 @@ module attributes { 
  spv.target_env = {

    version = 3 : i32,

    extensions = [],

    capabilities = [1: i32, 63: i32] // Shader, GroupNonUniformArithmetic

    capabilities = [1: i32, 63: i32], // Shader, GroupNonUniformArithmetic

    limits = {

      max_compute_workgroup_invocations = 128 : i32,

      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>

    }

  }

} {

func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>) attributes {
@@ -145,7 +157,11 @@ module attributes { 
  spv.target_env = {

    version = 3 : i32,

    extensions = [],

    capabilities = [1: i32, 63: i32] // Shader, GroupNonUniformArithmetic

    capabilities = [1: i32, 63: i32], // Shader, GroupNonUniformArithmetic

    limits = {

      max_compute_workgroup_invocations = 128 : i32,

      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>

    }

  }

} {

func @single_workgroup_reduction(%input: memref<16x8xi32>, %output: memref<16xi32>) attributes {