From 6d9a72ec80bbf4cfe10c81c944542ca195fc8d02 Mon Sep 17 00:00:00 2001
From: MaheshRavishankar <ravishankarm@google.com>
Date: Wed, 14 Oct 2020 22:32:52 -0700
Subject: [PATCH] [mlir][SPIRV] Adding an attribute to capture configuration
 for cooperative matrix operations.

Each hardware that supports SPV_C_CooperativeMatrixNV has a list of
configurations that are supported natively. Add an attribute to
specify the configurations supported to the `spv.target_env`.

Reviewed By: antiagainst, ThomasRaoux

Differential Revision: https://reviews.llvm.org/D89364
---
 mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td | 27 +++++++++++++++++-
 mlir/lib/Dialect/SPIRV/TargetAndABI.cpp         |  3 +-
 mlir/test/Dialect/SPIRV/target-and-abi.mlir     | 38 +++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td b/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td
index e8b1665..13bffab 100644
--- a/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td
+++ b/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td
@@ -36,6 +36,25 @@ def SPV_ExtensionArrayAttr : TypedArrayAttrBase<
 def SPV_CapabilityArrayAttr : TypedArrayAttrBase<
     SPV_CapabilityAttr, "SPIR-V capability array attribute">;
 
+// Description of cooperative matrix operations supported on the
+// target. Represents `VkCooperativeMatrixPropertiesNV`. See
+// https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkCooperativeMatrixPropertiesNV.html
+def SPV_CooperativeMatrixPropertiesNVAttr :
+    StructAttr<"CooperativeMatrixPropertiesNVAttr", SPIRV_Dialect, [
+    StructFieldAttr<"m_size", I32Attr>,
+    StructFieldAttr<"n_size", I32Attr>,
+    StructFieldAttr<"k_size", I32Attr>,
+    StructFieldAttr<"a_type", TypeAttr>,
+    StructFieldAttr<"b_type", TypeAttr>,
+    StructFieldAttr<"c_type", TypeAttr>,
+    StructFieldAttr<"result_type", TypeAttr>,
+    StructFieldAttr<"scope", SPV_ScopeAttr>
+]>;
+
+def SPV_CooperativeMatrixPropertiesNVArrayAttr :
+    TypedArrayAttrBase<SPV_CooperativeMatrixPropertiesNVAttr,
+                       "CooperativeMatrixPropertiesNV array attribute">;
+
 // This attribute specifies the limits for various resources on the target
 // architecture.
 //
@@ -60,7 +79,13 @@ def SPV_ResourceLimitsAttr : StructAttr<"ResourceLimitsAttr", SPIRV_Dialect, [
 
     // The default number of invocations in each subgroup.
     // 0x7FFFFFFF means unknown.
-    StructFieldAttr<"subgroup_size", DefaultValuedAttr<I32Attr, "0x7FFFFFFF">>
+    StructFieldAttr<"subgroup_size", DefaultValuedAttr<I32Attr, "0x7FFFFFFF">>,
+
+    // The configurations of cooperative matrix operations
+    // supported. Default is an empty list.
+    StructFieldAttr<
+        "cooperative_matrix_properties_nv",
+        DefaultValuedAttr<SPV_CooperativeMatrixPropertiesNVArrayAttr, "{}">>
 ]>;
 
 #endif // SPIRV_TARGET_AND_ABI
diff --git a/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp b/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp
index ae07651..29a89a3 100644
--- a/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp
+++ b/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp
@@ -140,7 +140,8 @@ spirv::getDefaultResourceLimits(MLIRContext *context) {
       /*max_compute_shared_memory_size=*/nullptr,
       /*max_compute_workgroup_invocations=*/nullptr,
       /*max_compute_workgroup_size=*/nullptr,
-      /*subgroup_size=*/nullptr, context);
+      /*subgroup_size=*/nullptr,
+      /*cooperative_matrix_properties_nv=*/nullptr, context);
 }
 
 StringRef spirv::getTargetEnvAttrName() { return "spv.target_env"; }
diff --git a/mlir/test/Dialect/SPIRV/target-and-abi.mlir b/mlir/test/Dialect/SPIRV/target-and-abi.mlir
index 6edc917..5f51680 100644
--- a/mlir/test/Dialect/SPIRV/target-and-abi.mlir
+++ b/mlir/test/Dialect/SPIRV/target-and-abi.mlir
@@ -171,6 +171,44 @@ func @target_env_extra_fields() attributes {
 
 // -----
 
+func @target_env_cooperative_matrix() attributes{
+  // CHECK:      spv.target_env = #spv.target_env<
+  // CHECK-SAME:   SPV_NV_cooperative_matrix
+  // CHECK-SAME:   cooperative_matrix_properties_nv = [
+  // CHECK-SAME:     {a_type = i8, b_type = i8, c_type = i32,
+  // CHECK-SAME:      k_size = 32 : i32, m_size = 8 : i32, n_size = 8 : i32
+  // CHECK-SAME:      result_type = i32, scope = 3 : i32}
+  // CHECK-SAME:     {a_type = f16, b_type = f16, c_type = f16,
+  // CHECK-SAME:      k_size = 16 : i32, m_size = 8 : i32, n_size = 8 : i32
+  // CHECK-SAME:      result_type = f16, scope = 3 : i32}
+  spv.target_env = #spv.target_env<
+  #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class,
+                            SPV_NV_cooperative_matrix]>,
+  {
+    cooperative_matrix_properties_nv = [{
+      m_size = 8: i32,
+      n_size = 8: i32,
+      k_size = 32: i32,
+      a_type = i8,
+      b_type = i8,
+      c_type = i32,
+      result_type = i32,
+      scope = 3: i32
+    }, {
+      m_size = 8: i32,
+      n_size = 8: i32,
+      k_size = 16: i32,
+      a_type = f16,
+      b_type = f16,
+      c_type = f16,
+      result_type = f16,
+      scope = 3: i32
+    }]
+  }>
+} { return }
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // spv.vce
 //===----------------------------------------------------------------------===//
-- 
2.7.4