spirv-opt: add support for tensors to type manager (KhronosGroup#6202)

As specified by SPV_ARM_tensors.


Change-Id: I229ea1a3232eb04a4124363d1641218acb298ac4

Signed-off-by: Kevin Petit <[email protected]>

Author: kpet

source/opt/type_manager.cpp

@@ -495,6 +495,36 @@ uint32_t TypeManager::GetTypeInstruction(const Type* type) {
               {SPV_OPERAND_TYPE_ID, {coop_vec->components()}}});
       break;
     }
+    case Type::kTensorARM: {
+      auto tensor_type = type->AsTensorARM();
+      uint32_t const element_type =
+          GetTypeInstruction(tensor_type->element_type());
+      if (element_type == 0) {
+        return 0;
+      }
+      if (tensor_type->rank_id() != 0) {
+        if (tensor_type->shape_id() != 0) {
+          typeInst = MakeUnique<Instruction>(
+              context(), spv::Op::OpTypeTensorARM, 0, id,
+              std::initializer_list<Operand>{
+                  {SPV_OPERAND_TYPE_ID, {element_type}},
+                  {SPV_OPERAND_TYPE_ID, {tensor_type->rank_id()}},
+                  {SPV_OPERAND_TYPE_ID, {tensor_type->shape_id()}}});
+        } else {
+          typeInst = MakeUnique<Instruction>(
+              context(), spv::Op::OpTypeTensorARM, 0, id,
+              std::initializer_list<Operand>{
+                  {SPV_OPERAND_TYPE_ID, {element_type}},
+                  {SPV_OPERAND_TYPE_ID, {tensor_type->rank_id()}}});
+        }
+      } else {
+        typeInst =
+            MakeUnique<Instruction>(context(), spv::Op::OpTypeTensorARM, 0, id,
+                                    std::initializer_list<Operand>{
+                                        {SPV_OPERAND_TYPE_ID, {element_type}}});
+      }
+      break;
+    }
     default:
       assert(false && "Unexpected type");
       break;
@@ -754,6 +784,14 @@ Type* TypeManager::RebuildType(uint32_t type_id, const Type& type) {
           cv_type->components());
       break;
     }
+    case Type::kTensorARM: {
+      const TensorARM* tensor_type = type.AsTensorARM();
+      const Type* element_type = tensor_type->element_type();
+      rebuilt_ty = MakeUnique<TensorARM>(
+          RebuildType(GetId(element_type), *element_type),
+          tensor_type->rank_id(), tensor_type->shape_id());
+      break;
+    }
     default:
       assert(false && "Unhandled type");
       return nullptr;
@@ -1036,6 +1074,23 @@ Type* TypeManager::RecordIfTypeDefinition(const Instruction& inst) {
                               inst.GetSingleWordInOperand(1), perm);
       break;
     }
+    case spv::Op::OpTypeTensorARM: {
+      switch (inst.NumInOperands()) {
+        case 1:
+          type = new TensorARM(GetType(inst.GetSingleWordInOperand(0)));
+          break;
+        case 2:
+          type = new TensorARM(GetType(inst.GetSingleWordInOperand(0)),
+                               inst.GetSingleWordInOperand(1));
+          break;
+        case 3:
+          type = new TensorARM(GetType(inst.GetSingleWordInOperand(0)),
+                               inst.GetSingleWordInOperand(1),
+                               inst.GetSingleWordInOperand(2));
+          break;
+      }
+      break;
+    }
     default:
       assert(false && "Type not handled by the type manager.");
       break;

source/opt/types.cpp

@@ -135,6 +135,7 @@ std::unique_ptr<Type> Type::Clone() const {
     DeclareKindCase(CooperativeVectorNV);
     DeclareKindCase(RayQueryKHR);
     DeclareKindCase(HitObjectNV);
+    DeclareKindCase(TensorARM);
 #undef DeclareKindCase
     default:
       assert(false && "Unhandled type");
@@ -187,6 +188,7 @@ bool Type::operator==(const Type& other) const {
     DeclareKindCase(HitObjectNV);
     DeclareKindCase(TensorLayoutNV);
     DeclareKindCase(TensorViewNV);
+    DeclareKindCase(TensorARM);
 #undef DeclareKindCase
     default:
       assert(false && "Unhandled type");
@@ -247,6 +249,7 @@ size_t Type::ComputeHashValue(size_t hash, SeenTypes* seen) const {
     DeclareKindCase(HitObjectNV);
     DeclareKindCase(TensorLayoutNV);
     DeclareKindCase(TensorViewNV);
+    DeclareKindCase(TensorARM);
 #undef DeclareKindCase
     default:
       assert(false && "Unhandled type");
@@ -899,6 +902,36 @@ bool CooperativeVectorNV::IsSameImpl(const Type* that,
          components_ == mt->components_ && HasSameDecorations(that);
 }
 
+TensorARM::TensorARM(const Type* elty, const uint32_t rank,
+                     const uint32_t shape)
+    : Type(kTensorARM), element_type_(elty), rank_id_(rank), shape_id_(shape) {
+  assert(elty != nullptr);
+  if (shape != 0) {
+    assert(rank != 0);
+  }
+}
+
+std::string TensorARM::str() const {
+  std::ostringstream oss;
+  oss << "tensor<" << element_type_->str() << ", id(" << rank_id_ << "), id("
+      << shape_id_ << ")>";
+  return oss.str();
+}
+
+size_t TensorARM::ComputeExtraStateHash(size_t hash, SeenTypes* seen) const {
+  hash = hash_combine(hash, rank_id_);
+  hash = hash_combine(hash, shape_id_);
+  return element_type_->ComputeHashValue(hash, seen);
+}
+
+bool TensorARM::IsSameImpl(const Type* that, IsSameCache* seen) const {
+  const TensorARM* tt = that->AsTensorARM();
+  if (!tt) return false;
+  return element_type_->IsSameImpl(tt->element_type_, seen) &&
+         rank_id_ == tt->rank_id_ && shape_id_ == tt->shape_id_ &&
+         HasSameDecorations(that);
+}
+
 }  // namespace analysis
 }  // namespace opt
 }  // namespace spvtools

source/opt/types.h

@@ -69,6 +69,7 @@ class RayQueryKHR;
 class HitObjectNV;
 class TensorLayoutNV;
 class TensorViewNV;
+class TensorARM;
 
 // Abstract class for a SPIR-V type. It has a bunch of As<sublcass>() methods,
 // which is used as a way to probe the actual <subclass>.
@@ -114,6 +115,7 @@ class Type {
     kHitObjectNV,
     kTensorLayoutNV,
     kTensorViewNV,
+    kTensorARM,
     kLast
   };
 
@@ -220,6 +222,7 @@ class Type {
   DeclareCastMethod(HitObjectNV)
   DeclareCastMethod(TensorLayoutNV)
   DeclareCastMethod(TensorViewNV)
+  DeclareCastMethod(TensorARM)
 #undef DeclareCastMethod
 
 protected:
@@ -774,6 +777,31 @@ class CooperativeVectorNV : public Type {
   const uint32_t components_;
 };
 
+class TensorARM : public Type {
+ public:
+  TensorARM(const Type* elty, const uint32_t rank = 0,
+            const uint32_t shape = 0);
+  TensorARM(const TensorARM&) = default;
+
+  std::string str() const override;
+
+  TensorARM* AsTensorARM() override { return this; }
+  const TensorARM* AsTensorARM() const override { return this; }
+
+  size_t ComputeExtraStateHash(size_t hash, SeenTypes* seen) const override;
+
+  const Type* element_type() const { return element_type_; }
+  uint32_t rank_id() const { return rank_id_; }
+  uint32_t shape_id() const { return shape_id_; }
+
+ private:
+  bool IsSameImpl(const Type* that, IsSameCache*) const override;
+
+  const Type* element_type_;
+  const uint32_t rank_id_;
+  const uint32_t shape_id_;
+};
+
 #define DefineParameterlessType(type, name)                                \
   class type : public Type {                                               \
    public:                                                                 \

test/opt/type_manager_test.cpp

@@ -182,6 +182,11 @@ std::vector<std::unique_ptr<Type>> GenerateAllTypes() {
   // SPV_AMDX_shader_enqueue
   types.emplace_back(new NodePayloadArrayAMDX(sts32f32));
 
+  // Tensors
+  types.emplace_back(new TensorARM(f32));
+  types.emplace_back(new TensorARM(f32, 4));
+  types.emplace_back(new TensorARM(f32, 4, 44));
+
   types.emplace_back(new TensorLayoutNV(1002, 1000));
   types.emplace_back(new TensorViewNV(1002, 1003, {1000, 1001}));
 
@@ -251,6 +256,11 @@ TEST(TypeManager, TypeStrings) {
     %id2    = OpConstant %u32 2
     %cmkhr  = OpTypeCooperativeMatrixKHR %f64 %id4 %id4 %id4 %id2
     %untyped = OpTypeUntypedPointerKHR Uniform
+    ; ID 43
+    %ts_shape = OpConstantComposite %a5u32 %id4 %id4 %id4 %id4
+    %ts  = OpTypeTensorARM %u32
+    %tsr = OpTypeTensorARM %u32 %id4
+    %tss = OpTypeTensorARM %u32 %id4 %ts_shape
   )";
 
   std::vector<std::pair<uint32_t, std::string>> type_id_strs = {
@@ -291,6 +301,10 @@ TEST(TypeManager, TypeStrings) {
       {39, "<float64, 6, 6, 6>"},
       {41, "<float64, 6, 6, 6, 40>"},
       {42, "untyped_ptr 2*"},  // Include storage class number
+      // Id 43 is OpConstantComposite %a5u32 %id4 %id4 %id4 %id4
+      {44, "tensor<uint32, id(0), id(0)>"},
+      {45, "tensor<uint32, id(6), id(0)>"},
+      {46, "tensor<uint32, id(6), id(43)>"},
   };
 
   std::unique_ptr<IRContext> context =
@@ -1049,8 +1063,11 @@ TEST(TypeManager, GetTypeInstructionAllTypes) {
 ; CHECK: [[uniform_ptr:%\w+]] = OpTypePointer Uniform [[uint]]
 ; CHECK: [[uint2:%\w+]] = OpConstant [[uint]] 2
 ; CHECK: [[uint8:%\w+]] = OpConstant [[uint]] 8
+; CHECK: [[uint4:%\w+]] = OpConstant [[uint]] 4
+; CHECK: [[uint_arr4:%\w+]] = OpTypeArray [[uint]] [[uint4]]
 ; CHECK: [[uint24:%\w+]] = OpConstant [[uint]] 24
 ; CHECK: [[uint42:%\w+]] = OpConstant [[uint]] 42
+; CHECK: [[uint_arr4_44:%\w+]] = OpConstantComposite [[uint_arr4]] [[uint4]] [[uint4]] [[uint4]] [[uint4]]
 ; CHECK: [[uint100:%\w+]] = OpConstant [[uint]] 100
 ; CHECK: [[void:%\w+]] = OpTypeVoid
 ; CHECK: [[bool:%\w+]] = OpTypeBool
@@ -1107,6 +1124,9 @@ TEST(TypeManager, GetTypeInstructionAllTypes) {
 ; CHECK: OpTypeCooperativeMatrixKHR [[f32]] [[uint8]] [[uint8]] [[uint8]] [[uint2]]
 ; CHECK: OpTypeRayQueryKHR
 ; CHECK: OpTypeHitObjectNV
+; CHECK: OpTypeTensorARM [[f32]]
+; CHECK: OpTypeTensorARM [[f32]] [[uint4]]
+; CHECK: OpTypeTensorARM [[f32]] [[uint4]] [[uint_arr4_44]]
 OpCapability Shader
 OpCapability Int64
 OpCapability Linkage
@@ -1118,8 +1138,11 @@ OpMemoryModel Logical GLSL450
 %1001 = OpConstant %uint 1
 %1002 = OpConstant %uint 2
 %8 = OpConstant %uint 8
+%4 = OpConstant %uint 4
+%5 = OpTypeArray %uint %4
 %24 = OpConstant %uint 24
 %42 = OpConstant %uint 42
+%44 = OpConstantComposite %5 %4 %4 %4 %4
 %100 = OpConstant %uint 100
 %1003 = OpConstantFalse %bool
   )";