Skip to content

Commit d8cc2f8

Browse files
kpetkpet
authored
spirv-opt: add support for tensor constants (KhronosGroup#6684)
Signed-off-by: Kevin Petit <kevin.petit@arm.com>
1 parent b3d6ad9 commit d8cc2f8

4 files changed

Lines changed: 409 additions & 1 deletion

File tree

‎source/opt/constants.cpp‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -287,6 +287,10 @@ std::unique_ptr<Constant> ConstantManager::CreateConstant(
287287
auto components = GetConstantsFromIds(literal_words_or_ids);
288288
if (components.empty()) return nullptr;
289289
return MakeUnique<ArrayConstant>(at, components);
290+
} else if (auto* tt = type->AsTensorARM()) {
291+
auto components = GetConstantsFromIds(literal_words_or_ids);
292+
if (components.empty()) return nullptr;
293+
return MakeUnique<TensorConstant>(tt, components);
290294
} else {
291295
return nullptr;
292296
}
@@ -376,6 +380,12 @@ std::unique_ptr<Instruction> ConstantManager::CreateCompositeInstruction(
376380
component_type_id = type_inst->GetSingleWordInOperand(component_index);
377381
} else if (type_inst && type_inst->opcode() == spv::Op::OpTypeArray) {
378382
component_type_id = type_inst->GetSingleWordInOperand(0);
383+
} else if (type_inst && type_inst->opcode() == spv::Op::OpTypeTensorARM) {
384+
component_type_id =
385+
context()->get_type_mgr()->GetId(component_const->type());
386+
if (component_type_id == 0) {
387+
return nullptr;
388+
}
379389
}
380390
uint32_t id = FindDeclaredConstant(component_const, component_type_id);
381391

@@ -431,6 +441,121 @@ const Constant* ConstantManager::GetNullCompositeConstant(const Type* type) {
431441
for (uint32_t i = 0; i < element_count; i++) {
432442
literal_words_or_id.push_back(null_id);
433443
}
444+
} else if (type->AsTensorARM()) {
445+
auto ttype = type->AsTensorARM();
446+
assert(ttype->is_shaped() && "Only shaped tensors are composites");
447+
const auto* shape_inst =
448+
context()->get_def_use_mgr()->GetDef(ttype->shape_id());
449+
if (!shape_inst) {
450+
return nullptr;
451+
}
452+
const auto* shape_cst = GetConstantFromInst(shape_inst);
453+
if (!shape_cst || !shape_cst->AsArrayConstant()) {
454+
return nullptr;
455+
}
456+
const auto& shape_components =
457+
shape_cst->AsArrayConstant()->GetComponents();
458+
if (shape_components.empty()) {
459+
return nullptr;
460+
}
461+
const auto* outer_dim_cst = shape_components.front()->AsIntConstant();
462+
if (!outer_dim_cst) {
463+
return nullptr;
464+
}
465+
const uint64_t element_count = outer_dim_cst->GetZeroExtendedValue();
466+
std::vector<const Constant*> components;
467+
components.reserve(static_cast<size_t>(element_count));
468+
469+
if (shape_components.size() == 1) {
470+
const Constant* element_null = GetConstant(ttype->element_type(), {});
471+
for (uint64_t i = 0; i < element_count; i++) {
472+
components.push_back(element_null);
473+
}
474+
return RegisterConstant(
475+
MakeUnique<TensorConstant>(ttype, std::move(components)));
476+
}
477+
478+
const Constant* rank_cst = nullptr;
479+
if (ttype->rank_id() != 0) {
480+
const auto* rank_inst =
481+
context()->get_def_use_mgr()->GetDef(ttype->rank_id());
482+
if (rank_inst) {
483+
rank_cst = GetConstantFromInst(rank_inst);
484+
}
485+
}
486+
uint64_t rank_value = shape_components.size();
487+
if (rank_cst && rank_cst->AsIntConstant()) {
488+
rank_value = rank_cst->AsIntConstant()->GetZeroExtendedValue();
489+
}
490+
if (rank_value <= 1) {
491+
const Constant* element_null = GetConstant(ttype->element_type(), {});
492+
for (uint64_t i = 0; i < element_count; i++) {
493+
components.push_back(element_null);
494+
}
495+
return RegisterConstant(
496+
MakeUnique<TensorConstant>(ttype, std::move(components)));
497+
}
498+
499+
const uint64_t inner_rank = rank_value - 1;
500+
const auto* rank_int_type =
501+
rank_cst ? rank_cst->type()->AsInteger() : nullptr;
502+
if (!rank_int_type) {
503+
rank_int_type = outer_dim_cst->type()->AsInteger();
504+
}
505+
if (!rank_int_type) {
506+
return nullptr;
507+
}
508+
const Constant* inner_rank_cst =
509+
GenerateIntegerConstant(rank_int_type, inner_rank);
510+
const uint32_t inner_rank_id =
511+
GetDefiningInstruction(inner_rank_cst)->result_id();
512+
513+
const Type* shape_elem_type = shape_components[1]->type();
514+
const auto* shape_elem_int = shape_elem_type->AsInteger();
515+
if (!shape_elem_int) {
516+
return nullptr;
517+
}
518+
const Constant* inner_shape_len_cst =
519+
GenerateIntegerConstant(shape_elem_int, inner_rank);
520+
const uint32_t inner_shape_len_id =
521+
GetDefiningInstruction(inner_shape_len_cst)->result_id();
522+
Array::LengthInfo inner_shape_len_info{
523+
inner_shape_len_id,
524+
{Array::LengthInfo::kConstant, static_cast<uint32_t>(inner_rank)}};
525+
Array inner_shape_type(shape_elem_type, inner_shape_len_info);
526+
const Type* inner_shape_reg_type =
527+
context()->get_type_mgr()->GetRegisteredType(&inner_shape_type);
528+
if (!inner_shape_reg_type) {
529+
return nullptr;
530+
}
531+
std::vector<uint32_t> inner_shape_ids;
532+
inner_shape_ids.reserve(shape_components.size() - 1);
533+
for (size_t i = 1; i < shape_components.size(); ++i) {
534+
inner_shape_ids.push_back(
535+
GetDefiningInstruction(shape_components[i])->result_id());
536+
}
537+
const Constant* inner_shape_cst =
538+
GetConstant(inner_shape_reg_type, inner_shape_ids);
539+
const uint32_t inner_shape_id =
540+
GetDefiningInstruction(inner_shape_cst)->result_id();
541+
542+
TensorARM inner_tensor_type(ttype->element_type(), inner_rank_id,
543+
inner_shape_id);
544+
const Type* inner_tensor_reg_type =
545+
context()->get_type_mgr()->GetRegisteredType(&inner_tensor_type);
546+
if (!inner_tensor_reg_type) {
547+
return nullptr;
548+
}
549+
const Constant* inner_null_tensor =
550+
GetNullCompositeConstant(inner_tensor_reg_type);
551+
if (!inner_null_tensor) {
552+
return nullptr;
553+
}
554+
for (uint64_t i = 0; i < element_count; i++) {
555+
components.push_back(inner_null_tensor);
556+
}
557+
return RegisterConstant(
558+
MakeUnique<TensorConstant>(ttype, std::move(components)));
434559
} else {
435560
return nullptr;
436561
}

‎source/opt/constants.h‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ class StructConstant;
5050
class VectorConstant;
5151
class MatrixConstant;
5252
class ArrayConstant;
53+
class TensorConstant;
5354
class NullConstant;
5455
class ConstantManager;
5556

@@ -73,6 +74,7 @@ class Constant {
7374
virtual VectorConstant* AsVectorConstant() { return nullptr; }
7475
virtual MatrixConstant* AsMatrixConstant() { return nullptr; }
7576
virtual ArrayConstant* AsArrayConstant() { return nullptr; }
77+
virtual TensorConstant* AsTensorConstant() { return nullptr; }
7678
virtual NullConstant* AsNullConstant() { return nullptr; }
7779

7880
virtual const ScalarConstant* AsScalarConstant() const { return nullptr; }
@@ -86,6 +88,7 @@ class Constant {
8688
virtual const VectorConstant* AsVectorConstant() const { return nullptr; }
8789
virtual const MatrixConstant* AsMatrixConstant() const { return nullptr; }
8890
virtual const ArrayConstant* AsArrayConstant() const { return nullptr; }
91+
virtual const TensorConstant* AsTensorConstant() const { return nullptr; }
8992
virtual const NullConstant* AsNullConstant() const { return nullptr; }
9093

9194
// Returns the float representation of the constant. Must be a 32 bit
@@ -426,6 +429,28 @@ class ArrayConstant : public CompositeConstant {
426429
}
427430
};
428431

432+
// Tensor type constant.
433+
class TensorConstant : public CompositeConstant {
434+
public:
435+
TensorConstant(const TensorARM* ty) : CompositeConstant(ty) {}
436+
TensorConstant(const TensorARM* ty,
437+
const std::vector<const Constant*>& components)
438+
: CompositeConstant(ty, components) {}
439+
TensorConstant(const TensorARM* ty, std::vector<const Constant*>&& components)
440+
: CompositeConstant(ty, std::move(components)) {}
441+
442+
TensorConstant* AsTensorConstant() override { return this; }
443+
const TensorConstant* AsTensorConstant() const override { return this; }
444+
445+
// Make a copy of this TensorConstant instance.
446+
std::unique_ptr<TensorConstant> CopyTensorConstant() const {
447+
return MakeUnique<TensorConstant>(type_->AsTensorARM(), components_);
448+
}
449+
std::unique_ptr<Constant> Copy() const override {
450+
return std::unique_ptr<Constant>(CopyTensorConstant().release());
451+
}
452+
};
453+
429454
// Null type constant.
430455
class NullConstant : public Constant {
431456
public:

‎test/opt/constant_manager_test.cpp‎

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,188 @@ OpMemoryModel Logical Simple
155155
}
156156
}
157157

158+
TEST_F(ConstantManagerTest, NullTensor) {
159+
const std::string text = R"(
160+
%uint = OpTypeInt 32 0
161+
%1 = OpConstant %uint 1
162+
%2 = OpConstant %uint 2
163+
%uint_5 = OpConstant %uint 5
164+
%uint_7 = OpConstant %uint 7
165+
%arr_uint_1 = OpTypeArray %uint %1
166+
%arr_uint_2 = OpTypeArray %uint %2
167+
%10 = OpConstantComposite %arr_uint_1 %uint_7
168+
%11 = OpConstantComposite %arr_uint_2 %uint_5 %uint_7
169+
)";
170+
171+
std::unique_ptr<IRContext> context =
172+
BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
173+
SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
174+
ASSERT_NE(context, nullptr);
175+
const auto cstmgr = context->get_constant_mgr();
176+
177+
Integer ty_uint(32, 0);
178+
179+
// rank-1 with 7 elements
180+
TensorARM ty_tensor_r1_7(&ty_uint, 1, 10);
181+
auto null_tensor_r1_7 =
182+
cstmgr->GetNullCompositeConstant(&ty_tensor_r1_7)->AsTensorConstant();
183+
ASSERT_NE(null_tensor_r1_7, nullptr);
184+
ASSERT_EQ(null_tensor_r1_7->GetComponents().size(), 7);
185+
ASSERT_TRUE(null_tensor_r1_7->IsZero());
186+
187+
// rank-2 with 5 elements of rank-1 type with 7 elements
188+
TensorARM ty_tensor_r2_5_7(&ty_uint, 2, 11);
189+
auto null_tensor_r2_5_7 =
190+
cstmgr->GetNullCompositeConstant(&ty_tensor_r2_5_7)->AsTensorConstant();
191+
ASSERT_NE(null_tensor_r2_5_7, nullptr);
192+
ASSERT_EQ(null_tensor_r2_5_7->GetComponents().size(), 5);
193+
ASSERT_TRUE(null_tensor_r2_5_7->IsZero());
194+
ASSERT_NE(null_tensor_r2_5_7->GetComponents()[0]->AsTensorConstant(),
195+
nullptr);
196+
ASSERT_TRUE(null_tensor_r2_5_7->GetComponents()[0]->type()->IsSameImpl(
197+
&ty_tensor_r1_7, nullptr));
198+
}
199+
200+
TEST_F(ConstantManagerTest, TensorConstantFromInstruction) {
201+
const std::string text = R"(
202+
OpCapability TensorsARM
203+
OpExtension "SPV_ARM_tensors"
204+
%1 = OpTypeInt 32 0
205+
%2 = OpConstant %1 1
206+
%3 = OpConstant %1 2
207+
%4 = OpConstant %1 3
208+
%5 = OpTypeArray %1 %2
209+
%6 = OpConstantComposite %5 %4
210+
%7 = OpTypeTensorARM %1 %2 %6
211+
%8 = OpConstantComposite %7 %2 %3 %4
212+
)";
213+
214+
std::unique_ptr<IRContext> context =
215+
BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
216+
SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
217+
ASSERT_NE(context, nullptr);
218+
219+
const Constant* constant =
220+
context->get_constant_mgr()->FindDeclaredConstant(8);
221+
ASSERT_NE(constant, nullptr);
222+
const auto* tensor_const = constant->AsTensorConstant();
223+
ASSERT_NE(tensor_const, nullptr);
224+
ASSERT_EQ(tensor_const->GetComponents().size(), 3);
225+
EXPECT_EQ(tensor_const->GetComponents()[0]->GetZeroExtendedValue(), 1u);
226+
EXPECT_EQ(tensor_const->GetComponents()[1]->GetZeroExtendedValue(), 2u);
227+
EXPECT_EQ(tensor_const->GetComponents()[2]->GetZeroExtendedValue(), 3u);
228+
EXPECT_FALSE(tensor_const->IsZero());
229+
}
230+
231+
TEST_F(ConstantManagerTest, NestedTensorConstantFromInstruction) {
232+
const std::string text = R"(
233+
OpCapability TensorsARM
234+
OpExtension "SPV_ARM_tensors"
235+
%1 = OpTypeInt 32 0
236+
%2 = OpConstant %1 1
237+
%3 = OpConstant %1 2
238+
%4 = OpConstant %1 3
239+
%5 = OpConstant %1 4
240+
%6 = OpConstant %1 5
241+
%7 = OpConstant %1 6
242+
%8 = OpTypeArray %1 %2
243+
%9 = OpTypeArray %1 %3
244+
%10 = OpConstantComposite %8 %4
245+
%11 = OpConstantComposite %9 %3 %4
246+
%12 = OpTypeTensorARM %1 %2 %10
247+
%13 = OpTypeTensorARM %1 %3 %11
248+
%14 = OpConstantComposite %12 %2 %3 %4
249+
%15 = OpConstantComposite %12 %5 %6 %7
250+
%16 = OpConstantComposite %13 %14 %15
251+
)";
252+
253+
std::unique_ptr<IRContext> context =
254+
BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
255+
SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
256+
ASSERT_NE(context, nullptr);
257+
258+
const Constant* constant =
259+
context->get_constant_mgr()->FindDeclaredConstant(16);
260+
ASSERT_NE(constant, nullptr);
261+
const auto* tensor_const = constant->AsTensorConstant();
262+
ASSERT_NE(tensor_const, nullptr);
263+
ASSERT_EQ(tensor_const->GetComponents().size(), 2);
264+
265+
const auto* row_0 = tensor_const->GetComponents()[0]->AsTensorConstant();
266+
const auto* row_1 = tensor_const->GetComponents()[1]->AsTensorConstant();
267+
ASSERT_NE(row_0, nullptr);
268+
ASSERT_NE(row_1, nullptr);
269+
ASSERT_EQ(row_0->GetComponents().size(), 3);
270+
ASSERT_EQ(row_1->GetComponents().size(), 3);
271+
EXPECT_EQ(row_0->GetComponents()[0]->GetZeroExtendedValue(), 1u);
272+
EXPECT_EQ(row_0->GetComponents()[1]->GetZeroExtendedValue(), 2u);
273+
EXPECT_EQ(row_0->GetComponents()[2]->GetZeroExtendedValue(), 3u);
274+
EXPECT_EQ(row_1->GetComponents()[0]->GetZeroExtendedValue(), 4u);
275+
EXPECT_EQ(row_1->GetComponents()[1]->GetZeroExtendedValue(), 5u);
276+
EXPECT_EQ(row_1->GetComponents()[2]->GetZeroExtendedValue(), 6u);
277+
EXPECT_FALSE(tensor_const->IsZero());
278+
}
279+
280+
TEST_F(ConstantManagerTest, GetDefiningInstructionForNestedTensorConstant) {
281+
const std::string text = R"(
282+
OpCapability TensorsARM
283+
OpExtension "SPV_ARM_tensors"
284+
%1 = OpTypeInt 32 0
285+
%2 = OpConstant %1 1
286+
%3 = OpConstant %1 2
287+
%4 = OpConstant %1 3
288+
%5 = OpConstant %1 4
289+
%6 = OpConstant %1 5
290+
%7 = OpConstant %1 6
291+
%8 = OpTypeArray %1 %2
292+
%9 = OpTypeArray %1 %3
293+
%10 = OpConstantComposite %8 %4
294+
%11 = OpConstantComposite %9 %3 %4
295+
%12 = OpTypeTensorARM %1 %2 %10
296+
%13 = OpTypeTensorARM %1 %3 %11
297+
)";
298+
299+
std::unique_ptr<IRContext> context =
300+
BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
301+
SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
302+
ASSERT_NE(context, nullptr);
303+
304+
ConstantManager* const_mgr = context->get_constant_mgr();
305+
Type* tensor_r1_type = context->get_type_mgr()->GetType(12);
306+
Type* tensor_r2_type = context->get_type_mgr()->GetType(13);
307+
ASSERT_NE(tensor_r1_type, nullptr);
308+
ASSERT_NE(tensor_r2_type, nullptr);
309+
310+
const Constant* row_0 = const_mgr->GetConstant(tensor_r1_type, {2, 3, 4});
311+
const Constant* row_1 = const_mgr->GetConstant(tensor_r1_type, {5, 6, 7});
312+
ASSERT_NE(row_0, nullptr);
313+
ASSERT_NE(row_1, nullptr);
314+
315+
Instruction* row_0_inst = const_mgr->GetDefiningInstruction(row_0, 12);
316+
Instruction* row_1_inst = const_mgr->GetDefiningInstruction(row_1, 12);
317+
ASSERT_NE(row_0_inst, nullptr);
318+
ASSERT_NE(row_1_inst, nullptr);
319+
EXPECT_EQ(row_0_inst->opcode(), spv::Op::OpConstantComposite);
320+
EXPECT_EQ(row_0_inst->type_id(), 12u);
321+
EXPECT_EQ(row_0_inst->NumInOperands(), 3u);
322+
EXPECT_EQ(row_0_inst->GetSingleWordInOperand(0), 2u);
323+
EXPECT_EQ(row_0_inst->GetSingleWordInOperand(1), 3u);
324+
EXPECT_EQ(row_0_inst->GetSingleWordInOperand(2), 4u);
325+
326+
const Constant* tensor_const = const_mgr->GetConstant(
327+
tensor_r2_type, {row_0_inst->result_id(), row_1_inst->result_id()});
328+
ASSERT_NE(tensor_const, nullptr);
329+
330+
Instruction* tensor_inst =
331+
const_mgr->GetDefiningInstruction(tensor_const, 13);
332+
ASSERT_NE(tensor_inst, nullptr);
333+
EXPECT_EQ(tensor_inst->opcode(), spv::Op::OpConstantComposite);
334+
EXPECT_EQ(tensor_inst->type_id(), 13u);
335+
EXPECT_EQ(tensor_inst->NumInOperands(), 2u);
336+
EXPECT_EQ(tensor_inst->GetSingleWordInOperand(0), row_0_inst->result_id());
337+
EXPECT_EQ(tensor_inst->GetSingleWordInOperand(1), row_1_inst->result_id());
338+
}
339+
158340
} // namespace
159341
} // namespace analysis
160342
} // namespace opt

0 commit comments

Comments
 (0)