Skip to content

Commit 86fd3af

Browse files
lhutton1lhutton1
authored
[mlir][tosa] Add ext-int64 support (llvm#164389)
This commit adds support for the EXT-INT64 extension added to the specification here: arm/tosa-specification@1b690f8
1 parent 332f786 commit 86fd3af

File tree

11 files changed

+236
-107
lines changed

11 files changed

+236
-107
lines changed

‎mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc‎

Lines changed: 127 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -467,6 +467,13 @@ profileComplianceMap = {
467467
extensionComplianceMap = {
468468
{"tosa.argmax",
469469
{{{Extension::int16}, {{{i16T, i32T}, SpecificationVersion::V_1_0}}},
470+
{{Extension::int64},
471+
{{{i8T, i64T}, SpecificationVersion::V_1_1_DRAFT},
472+
{{i16T, i64T}, SpecificationVersion::V_1_1_DRAFT},
473+
{{i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
474+
{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
475+
{{fp16T, i64T}, SpecificationVersion::V_1_1_DRAFT},
476+
{{fp32T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
470477
{{Extension::fp8e4m3}, {{{fp8e4m3T, i32T}, SpecificationVersion::V_1_0}}},
471478
{{Extension::fp8e5m2}, {{{fp8e5m2T, i32T}, SpecificationVersion::V_1_0}}},
472479
{{Extension::bf16}, {{{bf16T, i32T}, SpecificationVersion::V_1_0}}}}},
@@ -601,30 +608,68 @@ extensionComplianceMap = {
601608
{"tosa.tanh",
602609
{{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
603610
{"tosa.add",
604-
{{{Extension::bf16},
611+
{{{Extension::int64},
612+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
613+
{{Extension::bf16},
605614
{{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
615+
{"tosa.arithmetic_right_shift",
616+
{{{Extension::int64},
617+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
618+
{"tosa.bitwise_and",
619+
{{{Extension::int64},
620+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
621+
{"tosa.bitwise_or",
622+
{{{Extension::int64},
623+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
624+
{"tosa.bitwise_xor",
625+
{{{Extension::int64},
626+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
627+
{"tosa.intdiv",
628+
{{{Extension::int64},
629+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
630+
{"tosa.logical_left_shift",
631+
{{{Extension::int64},
632+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
633+
{"tosa.logical_right_shift",
634+
{{{Extension::int64},
635+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
606636
{"tosa.maximum",
607-
{{{Extension::bf16},
637+
{{{Extension::int64},
638+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
639+
{{Extension::bf16},
608640
{{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
609641
{"tosa.minimum",
610-
{{{Extension::bf16},
642+
{{{Extension::int64},
643+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
644+
{{Extension::bf16},
611645
{{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
612646
{"tosa.mul",
613-
{{{Extension::bf16},
647+
{{{Extension::int64},
648+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
649+
{{Extension::bf16},
614650
{{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
615651
{"tosa.pow",
616652
{{{Extension::bf16},
617653
{{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
618654
{"tosa.sub",
619-
{{{Extension::bf16},
655+
{{{Extension::int64},
656+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
657+
{{Extension::bf16},
620658
{{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
621659
{"tosa.table",
622660
{{{Extension::int16},
623661
{{{i16T, i16T, i32T}, SpecificationVersion::V_1_0}}}}},
624662
{"tosa.abs",
625-
{{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
663+
{{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
664+
{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
665+
{"tosa.bitwise_not",
666+
{{{Extension::int64},
667+
{{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
626668
{"tosa.ceil",
627669
{{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
670+
{"tosa.clz",
671+
{{{Extension::int64},
672+
{{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
628673
{"tosa.cos",
629674
{{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
630675
{"tosa.exp",
@@ -634,7 +679,9 @@ extensionComplianceMap = {
634679
{"tosa.log",
635680
{{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
636681
{"tosa.negate",
637-
{{{Extension::bf16},
682+
{{{Extension::int64},
683+
{{{i64T, i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
684+
{{Extension::bf16},
638685
{{{bf16T, bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
639686
{"tosa.reciprocal",
640687
{{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
@@ -643,65 +690,84 @@ extensionComplianceMap = {
643690
{"tosa.sin",
644691
{{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
645692
{"tosa.select",
646-
{{{Extension::bf16},
693+
{{{Extension::int64},
694+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
695+
{{Extension::bf16},
647696
{{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
648697
{"tosa.equal",
649-
{{{Extension::bf16},
698+
{{{Extension::int64},
699+
{{{i64T, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT}}},
700+
{{Extension::bf16},
650701
{{{bf16T, bf16T, boolT}, SpecificationVersion::V_1_0}}}}},
651702
{"tosa.greater",
652-
{{{Extension::bf16},
703+
{{{Extension::int64},
704+
{{{i64T, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT}}},
705+
{{Extension::bf16},
653706
{{{bf16T, bf16T, boolT}, SpecificationVersion::V_1_0}}}}},
654707
{"tosa.greater_equal",
655-
{{{Extension::bf16},
708+
{{{Extension::int64},
709+
{{{i64T, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT}}},
710+
{{Extension::bf16},
656711
{{{bf16T, bf16T, boolT}, SpecificationVersion::V_1_0}}}}},
657712
{"tosa.reduce_max",
658-
{{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
713+
{{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
714+
{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
659715
{"tosa.reduce_min",
660-
{{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
716+
{{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
717+
{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
661718
{"tosa.reduce_product",
662719
{{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
663720
{"tosa.reduce_sum",
664-
{{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
721+
{{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
722+
{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
665723
{"tosa.concat",
666724
{{{Extension::int16}, {{{i16T, i16T}, SpecificationVersion::V_1_0}}},
725+
{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
667726
{{Extension::fp8e4m3},
668727
{{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
669728
{{Extension::fp8e5m2},
670729
{{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
671730
{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
672731
{"tosa.pad",
673-
{{{Extension::fp8e4m3},
732+
{{{Extension::int64},
733+
{{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
734+
{{Extension::fp8e4m3},
674735
{{{fp8e4m3T, fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
675736
{{Extension::fp8e5m2},
676737
{{{fp8e5m2T, fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
677738
{{Extension::bf16},
678739
{{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
679740
{"tosa.reshape",
680-
{{{Extension::fp8e4m3},
741+
{{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
742+
{{Extension::fp8e4m3},
681743
{{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
682744
{{Extension::fp8e5m2},
683745
{{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
684746
{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
685747
{"tosa.reverse",
686-
{{{Extension::fp8e4m3},
748+
{{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
749+
{{Extension::fp8e4m3},
687750
{{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
688751
{{Extension::fp8e5m2},
689752
{{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
690753
{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
691754
{"tosa.slice",
692-
{{{Extension::fp8e4m3},
755+
{{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
756+
{{Extension::fp8e4m3},
693757
{{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
694758
{{Extension::fp8e5m2},
695759
{{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
696760
{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
697761
{"tosa.tile",
698-
{{{Extension::fp8e4m3},
762+
{{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
763+
{{Extension::fp8e4m3},
699764
{{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
700765
{{Extension::fp8e5m2},
701766
{{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
702767
{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
703768
{"tosa.transpose",
704-
{{{Extension::fp8e4m3},
769+
{{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
770+
{{Extension::fp8e4m3},
705771
{{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
706772
{{Extension::fp8e5m2},
707773
{{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
@@ -712,14 +778,48 @@ extensionComplianceMap = {
712778
{{Extension::fp8e5m2},
713779
{{{fp8e5m2T, i32T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
714780
{{Extension::bf16},
715-
{{{bf16T, i32T, bf16T}, SpecificationVersion::V_1_0}}}}},
781+
{{{bf16T, i32T, bf16T}, SpecificationVersion::V_1_0}}},
782+
{{Extension::int64},
783+
{{{i8T, i64T, i8T}, SpecificationVersion::V_1_1_DRAFT},
784+
{{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT},
785+
{{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT},
786+
{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
787+
{{fp16T, i64T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
788+
{{fp32T, i64T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}},
789+
{{Extension::fp8e4m3, Extension::int64},
790+
{{{fp8e4m3T, i64T, fp8e4m3T}, SpecificationVersion::V_1_1_DRAFT}},
791+
allOf},
792+
{{Extension::fp8e5m2, Extension::int64},
793+
{{{fp8e5m2T, i64T, fp8e5m2T}, SpecificationVersion::V_1_1_DRAFT}},
794+
allOf},
795+
{{Extension::bf16, Extension::int64},
796+
{{{bf16T, i64T, bf16T}, SpecificationVersion::V_1_1_DRAFT}},
797+
allOf}}},
716798
{"tosa.scatter",
717799
{{{Extension::fp8e4m3},
718800
{{{fp8e4m3T, i32T, fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
719801
{{Extension::fp8e5m2},
720802
{{{fp8e5m2T, i32T, fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
721803
{{Extension::bf16},
722-
{{{bf16T, i32T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
804+
{{{bf16T, i32T, bf16T, bf16T}, SpecificationVersion::V_1_0}}},
805+
{{Extension::int64},
806+
{{{i8T, i64T, i8T, i8T}, SpecificationVersion::V_1_1_DRAFT},
807+
{{i16T, i64T, i16T, i16T}, SpecificationVersion::V_1_1_DRAFT},
808+
{{i32T, i64T, i32T, i32T}, SpecificationVersion::V_1_1_DRAFT},
809+
{{i64T, i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
810+
{{fp16T, i64T, fp16T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
811+
{{fp32T, i64T, fp32T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}},
812+
{{Extension::fp8e4m3, Extension::int64},
813+
{{{fp8e4m3T, i64T, fp8e4m3T, fp8e4m3T},
814+
SpecificationVersion::V_1_1_DRAFT}},
815+
allOf},
816+
{{Extension::fp8e5m2, Extension::int64},
817+
{{{fp8e5m2T, i64T, fp8e5m2T, fp8e5m2T},
818+
SpecificationVersion::V_1_1_DRAFT}},
819+
allOf},
820+
{{Extension::bf16, Extension::int64},
821+
{{{bf16T, i64T, bf16T, bf16T}, SpecificationVersion::V_1_1_DRAFT}},
822+
allOf}}},
723823
{"tosa.resize",
724824
{{{Extension::int16},
725825
{{{i16T, i48T}, SpecificationVersion::V_1_0},
@@ -735,6 +835,9 @@ extensionComplianceMap = {
735835
{{bf16T, i32T}, SpecificationVersion::V_1_0},
736836
{{bf16T, fp32T}, SpecificationVersion::V_1_0},
737837
{{fp32T, bf16T}, SpecificationVersion::V_1_0}}},
838+
{{Extension::int64},
839+
{{{i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
840+
{{i64T, i32T}, SpecificationVersion::V_1_1_DRAFT}}},
738841
{{Extension::bf16, Extension::fp8e4m3},
739842
{{{bf16T, fp8e4m3T}, SpecificationVersion::V_1_0},
740843
{{fp8e4m3T, bf16T}, SpecificationVersion::V_1_0}},
@@ -761,12 +864,14 @@ extensionComplianceMap = {
761864
{"tosa.const",
762865
{{{Extension::int4}, {{{i4T}, SpecificationVersion::V_1_0}}},
763866
{{Extension::int16}, {{{i48T}, SpecificationVersion::V_1_0}}},
867+
{{Extension::int64}, {{{i64T}, SpecificationVersion::V_1_1_DRAFT}}},
764868
{{Extension::fp8e4m3}, {{{fp8e4m3T}, SpecificationVersion::V_1_0}}},
765869
{{Extension::fp8e5m2}, {{{fp8e5m2T}, SpecificationVersion::V_1_0}}},
766870
{{Extension::bf16}, {{{bf16T}, SpecificationVersion::V_1_0}}}}},
767871
{"tosa.identity",
768872
{{{Extension::int4}, {{{i4T, i4T}, SpecificationVersion::V_1_0}}},
769873
{{Extension::int16}, {{{i48T, i48T}, SpecificationVersion::V_1_0}}},
874+
{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
770875
{{Extension::fp8e4m3},
771876
{{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
772877
{{Extension::fp8e5m2},

‎mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,7 @@ class Tosa_I32EnumAttr<string name, string description, string mnemonic,
228228
// FP : Floating-Point Inference. Primarily FP16 and FP32 operations.
229229
//
230230
// Extension:
231+
// INT64 : 64-bit integer operations.
231232
// INT16 : 16-bit integer operations.
232233
// INT4 : 4-bit integer weights.
233234
// BF16 : BFloat16 operations.
@@ -271,21 +272,24 @@ def Tosa_EXT_DOUBLEROUND : I32EnumAttrCase<"doubleround", 9>;
271272
def Tosa_EXT_INEXACTROUND : I32EnumAttrCase<"inexactround", 10>;
272273
def Tosa_EXT_DYNAMIC : I32EnumAttrCase<"dynamic", 11>;
273274
def Tosa_EXT_MXFP : I32EnumAttrCase<"mxfp", 12>;
275+
def Tosa_EXT_INT64 : I32EnumAttrCase<"int64", 13>;
276+
274277

275278
def Tosa_ExtensionAttr
276279
: Tosa_I32EnumAttr<"Extension", "supported TOSA extensions", "ext", [
277280
Tosa_EXT_NONE, Tosa_EXT_INT16, Tosa_EXT_INT4, Tosa_EXT_BF16,
278281
Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_FFT, Tosa_EXT_VARIABLE,
279282
Tosa_EXT_CONTROLFLOW, Tosa_EXT_DOUBLEROUND, Tosa_EXT_INEXACTROUND,
280-
Tosa_EXT_DYNAMIC, Tosa_EXT_MXFP
283+
Tosa_EXT_DYNAMIC, Tosa_EXT_MXFP, Tosa_EXT_INT64
281284
]> {
282285
let extraClassDeclaration = [{
283-
static llvm::SmallVector<Extension, 11> getAllValues() {
286+
static llvm::SmallVector<Extension, 13> getAllValues() {
284287
return {
285288
Extension::int16, Extension::int4, Extension::bf16,
286289
Extension::fp8e4m3, Extension::fp8e5m2, Extension::fft,
287290
Extension::variable, Extension::controlflow, Extension::doubleround,
288-
Extension::inexactround, Extension::dynamic, Extension::mxfp
291+
Extension::inexactround, Extension::dynamic, Extension::mxfp,
292+
Extension::int64
289293
};
290294
}
291295
}];

0 commit comments

Comments
 (0)