spirv-opt: add support for OpTypeGraphARM to type manager (KhronosGroup#6247)

* spirv-opt: add support for OpTypeGraphARM to type manager

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

* improve one test

Change-Id: I3dedee7a790dfb89869bf8df6bb8e5b16a6f1c8d

---------

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

Author: kpet

source/opt/type_manager.cpp

@@ -525,6 +525,19 @@ uint32_t TypeManager::GetTypeInstruction(const Type* type) {
       }
       break;
     }
+    case Type::kGraphARM: {
+      auto const gty = type->AsGraphARM();
+      std::vector<Operand> ops;
+      ops.push_back(
+          Operand(SPV_OPERAND_TYPE_LITERAL_INTEGER, {gty->num_inputs()}));
+      for (auto iotype : gty->io_types()) {
+        uint32_t iotype_id = GetTypeInstruction(iotype);
+        ops.push_back(Operand(SPV_OPERAND_TYPE_ID, {iotype_id}));
+      }
+      typeInst = MakeUnique<Instruction>(context(), spv::Op::OpTypeGraphARM, 0,
+                                         id, ops);
+      break;
+    }
     default:
       assert(false && "Unexpected type");
       break;
@@ -792,6 +805,15 @@ Type* TypeManager::RebuildType(uint32_t type_id, const Type& type) {
           tensor_type->rank_id(), tensor_type->shape_id());
       break;
     }
+    case Type::kGraphARM: {
+      const GraphARM* graph_type = type.AsGraphARM();
+      std::vector<const Type*> io_types;
+      for (auto ioty : graph_type->io_types()) {
+        io_types.push_back(RebuildType(GetId(ioty), *ioty));
+      }
+      rebuilt_ty = MakeUnique<GraphARM>(graph_type->num_inputs(), io_types);
+      break;
+    }
     default:
       assert(false && "Unhandled type");
       return nullptr;
@@ -1091,6 +1113,14 @@ Type* TypeManager::RecordIfTypeDefinition(const Instruction& inst) {
       }
       break;
     }
+    case spv::Op::OpTypeGraphARM: {
+      std::vector<const Type*> io_types;
+      for (unsigned i = 1; i < inst.NumInOperands(); i++) {
+        io_types.push_back(GetType(inst.GetSingleWordInOperand(i)));
+      }
+      type = new GraphARM(inst.GetSingleWordInOperand(0), io_types);
+      break;
+    }
     default:
       assert(false && "Type not handled by the type manager.");
       break;

source/opt/types.cpp

@@ -136,6 +136,7 @@ std::unique_ptr<Type> Type::Clone() const {
     DeclareKindCase(RayQueryKHR);
     DeclareKindCase(HitObjectNV);
     DeclareKindCase(TensorARM);
+    DeclareKindCase(GraphARM);
 #undef DeclareKindCase
     default:
       assert(false && "Unhandled type");
@@ -189,6 +190,7 @@ bool Type::operator==(const Type& other) const {
     DeclareKindCase(TensorLayoutNV);
     DeclareKindCase(TensorViewNV);
     DeclareKindCase(TensorARM);
+    DeclareKindCase(GraphARM);
 #undef DeclareKindCase
     default:
       assert(false && "Unhandled type");
@@ -250,6 +252,7 @@ size_t Type::ComputeHashValue(size_t hash, SeenTypes* seen) const {
     DeclareKindCase(TensorLayoutNV);
     DeclareKindCase(TensorViewNV);
     DeclareKindCase(TensorARM);
+    DeclareKindCase(GraphARM);
 #undef DeclareKindCase
     default:
       assert(false && "Unhandled type");
@@ -932,6 +935,61 @@ bool TensorARM::IsSameImpl(const Type* that, IsSameCache* seen) const {
          HasSameDecorations(that);
 }
 
+GraphARM::GraphARM(const uint32_t num_inputs,
+                   const std::vector<const Type*>& io_types)
+    : Type(kGraphARM), num_inputs_(num_inputs), io_types_(io_types) {
+  assert(io_types.size() > 0);
+}
+
+std::string GraphARM::str() const {
+  std::ostringstream oss;
+  oss << "graph<" << num_inputs_;
+  for (auto ioty : io_types_) {
+    oss << "," << ioty->str();
+  }
+  oss << ">";
+  return oss.str();
+}
+
+bool GraphARM::is_shaped() const {
+  // A graph is considered to be shaped if all its interface tensors are shaped
+  for (auto ioty : io_types_) {
+    auto tensor_type = ioty->AsTensorARM();
+    assert(tensor_type);
+    if (!tensor_type->is_shaped()) {
+      return false;
+    }
+  }
+  return true;
+}
+
+size_t GraphARM::ComputeExtraStateHash(size_t hash, SeenTypes* seen) const {
+  hash = hash_combine(hash, num_inputs_);
+  for (auto ioty : io_types_) {
+    hash = ioty->ComputeHashValue(hash, seen);
+  }
+  return hash;
+}
+
+bool GraphARM::IsSameImpl(const Type* that, IsSameCache* seen) const {
+  const GraphARM* og = that->AsGraphARM();
+  if (!og) {
+    return false;
+  }
+  if (num_inputs_ != og->num_inputs_) {
+    return false;
+  }
+  if (io_types_.size() != og->io_types_.size()) {
+    return false;
+  }
+  for (size_t i = 0; i < io_types_.size(); i++) {
+    if (!io_types_[i]->IsSameImpl(og->io_types_[i], seen)) {
+      return false;
+    }
+  }
+  return true;
+}
+
 }  // namespace analysis
 }  // namespace opt
 }  // namespace spvtools

source/opt/types.h

@@ -70,6 +70,7 @@ class HitObjectNV;
 class TensorLayoutNV;
 class TensorViewNV;
 class TensorARM;
+class GraphARM;
 
 // 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>.
@@ -116,6 +117,7 @@ class Type {
     kTensorLayoutNV,
     kTensorViewNV,
     kTensorARM,
+    kGraphARM,
     kLast
   };
 
@@ -223,6 +225,7 @@ class Type {
   DeclareCastMethod(TensorLayoutNV)
   DeclareCastMethod(TensorViewNV)
   DeclareCastMethod(TensorARM)
+  DeclareCastMethod(GraphARM)
 #undef DeclareCastMethod
 
 protected:
@@ -793,6 +796,8 @@ class TensorARM : public Type {
   const Type* element_type() const { return element_type_; }
   uint32_t rank_id() const { return rank_id_; }
   uint32_t shape_id() const { return shape_id_; }
+  bool is_ranked() const { return rank_id_ != 0; }
+  bool is_shaped() const { return shape_id_ != 0; }
 
  private:
   bool IsSameImpl(const Type* that, IsSameCache*) const override;
@@ -802,6 +807,29 @@ class TensorARM : public Type {
   const uint32_t shape_id_;
 };
 
+class GraphARM : public Type {
+ public:
+  GraphARM(const uint32_t num_inputs, const std::vector<const Type*>& io_types);
+  GraphARM(const GraphARM&) = default;
+
+  std::string str() const override;
+
+  GraphARM* AsGraphARM() override { return this; }
+  const GraphARM* AsGraphARM() const override { return this; }
+
+  uint32_t num_inputs() const { return num_inputs_; }
+  const std::vector<const Type*>& io_types() const { return io_types_; }
+  bool is_shaped() const;
+
+  size_t ComputeExtraStateHash(size_t hash, SeenTypes* seen) const override;
+
+ private:
+  bool IsSameImpl(const Type* that, IsSameCache*) const override;
+
+  const uint32_t num_inputs_;
+  const std::vector<const Type*> io_types_;
+};
+
 #define DefineParameterlessType(type, name)                                \
   class type : public Type {                                               \
    public:                                                                 \

test/opt/type_manager_test.cpp

@@ -184,8 +184,16 @@ std::vector<std::unique_ptr<Type>> GenerateAllTypes() {
 
   // Tensors
   types.emplace_back(new TensorARM(f32));
+  auto* tensor_f32 = types.back().get();
   types.emplace_back(new TensorARM(f32, 4));
+  auto* tensor_f32_ranked = types.back().get();
   types.emplace_back(new TensorARM(f32, 4, 44));
+  auto* tensor_f32_shaped = types.back().get();
+
+  // Graph
+  types.emplace_back(new GraphARM(0, {tensor_f32}));
+  types.emplace_back(new GraphARM(1, {tensor_f32_ranked, tensor_f32_ranked}));
+  types.emplace_back(new GraphARM(1, {tensor_f32_shaped, tensor_f32_shaped}));
 
   types.emplace_back(new TensorLayoutNV(1002, 1000));
   types.emplace_back(new TensorViewNV(1002, 1003, {1000, 1001}));
@@ -261,6 +269,9 @@ TEST(TypeManager, TypeStrings) {
     %ts  = OpTypeTensorARM %u32
     %tsr = OpTypeTensorARM %u32 %id4
     %tss = OpTypeTensorARM %u32 %id4 %ts_shape
+    %g_noin = OpTypeGraphARM 0 %ts
+    %g_onein = OpTypeGraphARM 1 %tsr %tsr
+    %g_shaped = OpTypeGraphARM 1 %tss %tss
   )";
 
   std::vector<std::pair<uint32_t, std::string>> type_id_strs = {
@@ -305,6 +316,11 @@ TEST(TypeManager, TypeStrings) {
       {44, "tensor<uint32, id(0), id(0)>"},
       {45, "tensor<uint32, id(6), id(0)>"},
       {46, "tensor<uint32, id(6), id(43)>"},
+      {47, "graph<0,tensor<uint32, id(0), id(0)>>"},
+      {48,
+       "graph<1,tensor<uint32, id(6), id(0)>,tensor<uint32, id(6), id(0)>>"},
+      {49,
+       "graph<1,tensor<uint32, id(6), id(43)>,tensor<uint32, id(6), id(43)>>"},
   };
 
   std::unique_ptr<IRContext> context =
@@ -1124,9 +1140,12 @@ 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]]
+; CHECK: [[tensor_f32:%\w+]] = OpTypeTensorARM [[f32]]
+; CHECK: [[tensor_f32_ranked:%\w+]] = OpTypeTensorARM [[f32]] [[uint4]]
+; CHECK: [[tensor_f32_shaped:%\w+]] = OpTypeTensorARM [[f32]] [[uint4]] [[uint_arr4_44]]
+; CHECK: OpTypeGraphARM 0 [[tensor_f32]]
+; CHECK: OpTypeGraphARM 1 [[tensor_f32_ranked]] [[tensor_f32_ranked]]
+; CHECK: OpTypeGraphARM 1 [[tensor_f32_shaped]] [[tensor_f32_shaped]]
 OpCapability Shader
 OpCapability Int64
 OpCapability Linkage