Skip to content

[mlir][tosa] Add ext-int64 support #164389

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 2 commits into from
Oct 24, 2025
Merged

[mlir][tosa] Add ext-int64 support #164389

merged 2 commits into from
Oct 24, 2025

Conversation

@lhutton1
Copy link
Contributor

@lhutton1 lhutton1 commented Oct 21, 2025 •
edited
Loading

This commit adds support for the EXT-INT64 extension added
to the specification here: arm/tosa-specification@1b690f8

Note: this change is dependent on other open pull requests such as #163433

@psunn psunn requested review from Jerry-Ge and psunn October 23, 2025 12:51
@lhutton1
[mlir][tosa] Add ext-int64 support
4f5d6e8 
This commit adds support for the EXT-INT64 extension added
to the specification here: https://review.mlplatform.org/c/tosa/specification/+/15365

Change-Id: Ibd7041492922587df951f0b600103af2d0102860
@lhutton1 lhutton1 marked this pull request as ready for review October 23, 2025 15:53
@lhutton1
Copy link
Contributor Author

This has been rebased & is ready for review.

@llvmbot
Copy link
Member

llvmbot commented Oct 23, 2025 •
edited
Loading

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-tosa

Author: Luke Hutton (lhutton1)

Changes

This commit adds support for the EXT-INT64 extension added
to the specification here: arm/tosa-specification@1b690f8

Note: this change is dependent on other open pull requests such as #163433

Patch is 50.27 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/164389.diff

11 Files Affected:

  • (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc (+127-22)
  • (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td (+6-3)
  • (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td (+36-36)
  • (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaProfileCompliance.h (+1)
  • (modified) mlir/lib/Dialect/Tosa/IR/TargetEnv.cpp (+1)
  • (modified) mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp (+1)
  • (modified) mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp (+1)
  • (modified) mlir/test/Dialect/Tosa/availability.mlir (+36-36)
  • (modified) mlir/test/Dialect/Tosa/invalid.mlir (-9)
  • (modified) mlir/test/Dialect/Tosa/invalid_extension.mlir (+9)
  • (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir (+17-1)
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
index 294fb9d99fdb6..06e4ee0c4176d 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
@@ -467,6 +467,13 @@ profileComplianceMap = {
 extensionComplianceMap = {
     {"tosa.argmax",
      {{{Extension::int16}, {{{i16T, i32T}, SpecificationVersion::V_1_0}}},
+      {{Extension::int64},
+       {{{i8T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i16T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp16T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp32T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
       {{Extension::fp8e4m3}, {{{fp8e4m3T, i32T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2}, {{{fp8e5m2T, i32T}, SpecificationVersion::V_1_0}}},
       {{Extension::bf16}, {{{bf16T, i32T}, SpecificationVersion::V_1_0}}}}},
@@ -601,30 +608,68 @@ extensionComplianceMap = {
     {"tosa.tanh",
      {{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.add",
-     {{{Extension::bf16},
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16},
        {{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
+    {"tosa.arithmetic_right_shift",
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
+    {"tosa.bitwise_and",
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
+    {"tosa.bitwise_or",
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
+    {"tosa.bitwise_xor",
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
+    {"tosa.intdiv",
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
+    {"tosa.logical_left_shift",
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
+    {"tosa.logical_right_shift",
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
     {"tosa.maximum",
-     {{{Extension::bf16},
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16},
        {{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.minimum",
-     {{{Extension::bf16},
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16},
        {{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.mul",
-     {{{Extension::bf16},
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16},
        {{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.pow",
      {{{Extension::bf16},
        {{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.sub",
-     {{{Extension::bf16},
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16},
        {{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.table",
      {{{Extension::int16},
        {{{i16T, i16T, i32T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.abs",
-     {{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
+     {{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
+    {"tosa.bitwise_not",
+     {{{Extension::int64},
+       {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
     {"tosa.ceil",
      {{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
+    {"tosa.clz",
+     {{{Extension::int64},
+       {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}}}},
     {"tosa.cos",
      {{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.exp",
@@ -634,7 +679,9 @@ extensionComplianceMap = {
     {"tosa.log",
      {{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.negate",
-     {{{Extension::bf16},
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16},
        {{{bf16T, bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.reciprocal",
      {{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
@@ -643,65 +690,84 @@ extensionComplianceMap = {
     {"tosa.sin",
      {{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.select",
-     {{{Extension::bf16},
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16},
        {{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.equal",
-     {{{Extension::bf16},
+     {{{Extension::int64},
+       {{{i64T, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16},
        {{{bf16T, bf16T, boolT}, SpecificationVersion::V_1_0}}}}},
     {"tosa.greater",
-     {{{Extension::bf16},
+     {{{Extension::int64},
+       {{{i64T, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16},
        {{{bf16T, bf16T, boolT}, SpecificationVersion::V_1_0}}}}},
     {"tosa.greater_equal",
-     {{{Extension::bf16},
+     {{{Extension::int64},
+       {{{i64T, i64T, boolT}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16},
        {{{bf16T, bf16T, boolT}, SpecificationVersion::V_1_0}}}}},
     {"tosa.reduce_max",
-     {{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
+     {{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.reduce_min",
-     {{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
+     {{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.reduce_product",
      {{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.reduce_sum",
-     {{{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
+     {{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.concat",
      {{{Extension::int16}, {{{i16T, i16T}, SpecificationVersion::V_1_0}}},
+      {{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
       {{Extension::fp8e4m3},
        {{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2},
        {{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
       {{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.pad",
-     {{{Extension::fp8e4m3},
+     {{{Extension::int64},
+       {{{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::fp8e4m3},
        {{{fp8e4m3T, fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2},
        {{{fp8e5m2T, fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
       {{Extension::bf16},
        {{{bf16T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.reshape",
-     {{{Extension::fp8e4m3},
+     {{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::fp8e4m3},
        {{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2},
        {{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
       {{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.reverse",
-     {{{Extension::fp8e4m3},
+     {{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::fp8e4m3},
        {{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2},
        {{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
       {{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.slice",
-     {{{Extension::fp8e4m3},
+     {{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::fp8e4m3},
        {{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2},
        {{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
       {{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.tile",
-     {{{Extension::fp8e4m3},
+     {{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::fp8e4m3},
        {{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2},
        {{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
       {{Extension::bf16}, {{{bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.transpose",
-     {{{Extension::fp8e4m3},
+     {{{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::fp8e4m3},
        {{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2},
        {{{fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
@@ -712,14 +778,48 @@ extensionComplianceMap = {
       {{Extension::fp8e5m2},
        {{{fp8e5m2T, i32T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
       {{Extension::bf16},
-       {{{bf16T, i32T, bf16T}, SpecificationVersion::V_1_0}}}}},
+       {{{bf16T, i32T, bf16T}, SpecificationVersion::V_1_0}}},
+      {{Extension::int64},
+       {{{i8T, i64T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp16T, i64T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp32T, i64T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::fp8e4m3, Extension::int64},
+       {{{fp8e4m3T, i64T, fp8e4m3T}, SpecificationVersion::V_1_1_DRAFT}},
+       allOf},
+      {{Extension::fp8e5m2, Extension::int64},
+       {{{fp8e5m2T, i64T, fp8e5m2T}, SpecificationVersion::V_1_1_DRAFT}},
+       allOf},
+      {{Extension::bf16, Extension::int64},
+       {{{bf16T, i64T, bf16T}, SpecificationVersion::V_1_1_DRAFT}},
+       allOf}}},
     {"tosa.scatter",
      {{{Extension::fp8e4m3},
        {{{fp8e4m3T, i32T, fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2},
        {{{fp8e5m2T, i32T, fp8e5m2T, fp8e5m2T}, SpecificationVersion::V_1_0}}},
       {{Extension::bf16},
-       {{{bf16T, i32T, bf16T, bf16T}, SpecificationVersion::V_1_0}}}}},
+       {{{bf16T, i32T, bf16T, bf16T}, SpecificationVersion::V_1_0}}},
+      {{Extension::int64},
+       {{{i8T, i64T, i8T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i16T, i64T, i16T, i16T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i32T, i64T, i32T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i64T, i64T, i64T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp16T, i64T, fp16T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp32T, i64T, fp32T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}},
+      {{Extension::fp8e4m3, Extension::int64},
+       {{{fp8e4m3T, i64T, fp8e4m3T, fp8e4m3T},
+         SpecificationVersion::V_1_1_DRAFT}},
+       allOf},
+      {{Extension::fp8e5m2, Extension::int64},
+       {{{fp8e5m2T, i64T, fp8e5m2T, fp8e5m2T},
+         SpecificationVersion::V_1_1_DRAFT}},
+       allOf},
+      {{Extension::bf16, Extension::int64},
+       {{{bf16T, i64T, bf16T, bf16T}, SpecificationVersion::V_1_1_DRAFT}},
+       allOf}}},
     {"tosa.resize",
      {{{Extension::int16},
        {{{i16T, i48T}, SpecificationVersion::V_1_0},
@@ -735,6 +835,9 @@ extensionComplianceMap = {
         {{bf16T, i32T}, SpecificationVersion::V_1_0},
         {{bf16T, fp32T}, SpecificationVersion::V_1_0},
         {{fp32T, bf16T}, SpecificationVersion::V_1_0}}},
+      {{Extension::int64},
+       {{{i32T, i64T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i64T, i32T}, SpecificationVersion::V_1_1_DRAFT}}},
       {{Extension::bf16, Extension::fp8e4m3},
        {{{bf16T, fp8e4m3T}, SpecificationVersion::V_1_0},
         {{fp8e4m3T, bf16T}, SpecificationVersion::V_1_0}},
@@ -761,12 +864,14 @@ extensionComplianceMap = {
     {"tosa.const",
      {{{Extension::int4}, {{{i4T}, SpecificationVersion::V_1_0}}},
       {{Extension::int16}, {{{i48T}, SpecificationVersion::V_1_0}}},
+      {{Extension::int64}, {{{i64T}, SpecificationVersion::V_1_1_DRAFT}}},
       {{Extension::fp8e4m3}, {{{fp8e4m3T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2}, {{{fp8e5m2T}, SpecificationVersion::V_1_0}}},
       {{Extension::bf16}, {{{bf16T}, SpecificationVersion::V_1_0}}}}},
     {"tosa.identity",
      {{{Extension::int4}, {{{i4T, i4T}, SpecificationVersion::V_1_0}}},
       {{Extension::int16}, {{{i48T, i48T}, SpecificationVersion::V_1_0}}},
+      {{Extension::int64}, {{{i64T, i64T}, SpecificationVersion::V_1_1_DRAFT}}},
       {{Extension::fp8e4m3},
        {{{fp8e4m3T, fp8e4m3T}, SpecificationVersion::V_1_0}}},
       {{Extension::fp8e5m2},
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
index 48e0073c76ab6..3adcfc8e3bd23 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
@@ -271,21 +271,24 @@ def Tosa_EXT_DOUBLEROUND  : I32EnumAttrCase<"doubleround", 9>;
 def Tosa_EXT_INEXACTROUND : I32EnumAttrCase<"inexactround", 10>;
 def Tosa_EXT_DYNAMIC      : I32EnumAttrCase<"dynamic", 11>;
 def Tosa_EXT_MXFP         : I32EnumAttrCase<"mxfp", 12>;
+def Tosa_EXT_INT64        : I32EnumAttrCase<"int64", 13>;
+
 
 def Tosa_ExtensionAttr
     : Tosa_I32EnumAttr<"Extension", "supported TOSA extensions", "ext", [
       Tosa_EXT_NONE, Tosa_EXT_INT16, Tosa_EXT_INT4, Tosa_EXT_BF16,
       Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_FFT, Tosa_EXT_VARIABLE,
       Tosa_EXT_CONTROLFLOW, Tosa_EXT_DOUBLEROUND, Tosa_EXT_INEXACTROUND,
-      Tosa_EXT_DYNAMIC, Tosa_EXT_MXFP
+      Tosa_EXT_DYNAMIC, Tosa_EXT_MXFP, Tosa_EXT_INT64
     ]> {
   let extraClassDeclaration = [{
-    static llvm::SmallVector<Extension, 11> getAllValues() {
+    static llvm::SmallVector<Extension, 14> getAllValues() {
       return {
         Extension::int16, Extension::int4, Extension::bf16,
         Extension::fp8e4m3, Extension::fp8e5m2, Extension::fft,
         Extension::variable, Extension::controlflow, Extension::doubleround,
-        Extension::inexactround, Extension::dynamic, Extension::mxfp
+        Extension::inexactround, Extension::dynamic, Extension::mxfp,
+        Extension::int64
       };
     }
   }];
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 6f07247b478c8..ddd8cba5f9dd5 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -52,7 +52,7 @@ def Tosa_ArgMaxOp : Tosa_InferShapedTypeOp<"argmax"> {
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[Tosa_EXT_INT16, Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16]>,
+    Extension<[Tosa_EXT_INT16, Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16, Tosa_EXT_INT64]>,
   ];
 
   let hasFolder = 1;
@@ -681,7 +681,7 @@ def Tosa_AddOp : Tosa_ElementwiseOp<"add", [
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[Tosa_EXT_BF16]>,
+    Extension<[Tosa_EXT_BF16, Tosa_EXT_INT64]>,
   ];
 
   let hasFolder = 1;
@@ -714,7 +714,7 @@ def Tosa_ArithmeticRightShiftOp : Tosa_ElementwiseOp<"arithmetic_right_shift",
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT]>,
-    Extension<[]>,
+    Extension<[Tosa_EXT_INT64]>,
   ];
 
   let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
@@ -744,7 +744,7 @@ def Tosa_BitwiseAndOp : Tosa_ElementwiseOp<"bitwise_and", [
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT]>,
-    Extension<[]>,
+    Extension<[Tosa_EXT_INT64]>,
   ];
 
   let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
@@ -774,7 +774,7 @@ def Tosa_BitwiseOrOp : Tosa_ElementwiseOp<"bitwise_or", [
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT]>,
-    Extension<[]>,
+    Extension<[Tosa_EXT_INT64]>,
   ];
 
   let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
@@ -804,7 +804,7 @@ def Tosa_BitwiseXorOp : Tosa_ElementwiseOp<"bitwise_xor", [
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT]>,
-    Extension<[]>,
+    Extension<[Tosa_EXT_INT64]>,
   ];
 
   let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
@@ -835,7 +835,7 @@ def Tosa_IntDivOp : Tosa_ElementwiseOp<"intdiv", [SameOperandsAndResultElementTy
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[]>,
+    Extension<[Tosa_EXT_INT64]>,
   ];
 
   let hasFolder = 1;
@@ -897,7 +897,7 @@ def Tosa_LogicalLeftShiftOp : Tosa_ElementwiseOp<"logical_left_shift",
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[]>,
+    Extension<[Tosa_EXT_INT64]>,
   ];
 
   let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
@@ -927,7 +927,7 @@ def Tosa_LogicalRightShiftOp : Tosa_ElementwiseOp<"logical_right_shift",
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[]>,
+    Extension<[Tosa_EXT_INT64]>,
   ];
 
   let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
@@ -1018,7 +1018,7 @@ def Tosa_MaximumOp : Tosa_ElementwiseOp<"maximum", [
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[Tosa_EXT_BF16]>,
+    Extension<[Tosa_EXT_BF16, Tosa_EXT_INT64]>,
   ];
   let hasCustomAssemblyFormat = 1;
 }
@@ -1048,7 +1048,7 @@ def Tosa_MinimumOp : Tosa_ElementwiseOp<"minimum", [
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[Tosa_EXT_BF16]>,
+    Extension<[Tosa_EXT_BF16, Tosa_EXT_INT64]>,
   ];
 
   let hasCustomAssemblyFormat = 1;
@@ -1082,7 +1082,7 @@ def Tosa_MulOp : Tosa_Op<"mul", [
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[Tosa_EXT_BF16]>,
+    Extension<[Tosa_EXT_BF16, Tosa_EXT_INT64]>,
   ];
 
   let hasFolder = 1;
@@ -1143,7 +1143,7 @@ def Tosa_SubOp : Tosa_ElementwiseOp<"sub", [SameOperandsAndResultElementType]> {
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[Tosa_EXT_BF16]>,
+    Extension<[Tosa_EXT_BF16, Tosa_EXT_INT64]>,
   ];
 
   let hasFolder = 1;
@@ -1228,7 +1228,7 @@ def Tosa_AbsOp : Tosa_ElementwiseUnaryOp<"abs"> {
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[Tosa_EXT_BF16]>,
+    Extension<[Tosa_EXT_BF16, Tosa_EXT_INT64]>,
   ];
 
   let hasFolder = 1;
@@ -1256,7 +1256,7 @@ def Tosa_BitwiseNotOp : Tosa_ElementwiseUnaryOp<"bitwise_not"> {
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT]>,
-    Extension<[]>,
+    Extension<[Tosa_EXT_INT64]>,
   ];
 
   let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
@@ -1308,7 +1308,7 @@ def Tosa_ClzOp : Tosa_ElementwiseUnaryOp<"clz"> {
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT]>,
-    Extension<[]>,
+    Extension<[Tosa_EXT_INT64]>,
   ];
 
   let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
@@ -1468,7 +1468,7 @@ def Tosa_NegateOp : Tosa_InferShapedTypeOp<"negate", [
 
   list<Availability> availability = [
     Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
-    Extension<[Tosa_EXT_BF16]>,
+    Extension<[Tosa_EXT_BF16, Tosa_EXT_INT64]>,
   ];
 
   let builders = [Tosa_NegateOpQuantInfoBuilder];
@@ -1604,7 +1604,7 @@ def To...
[truncated]

@psunn
Copy link
Contributor

psunn commented Oct 23, 2025

Thank you for updating this PR. I’ll review it later today.

psunn
psunn requested changes Oct 23, 2025
View reviewed changes
Copy link
Contributor

@psunn psunn left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good overall—thanks!
I have one minor change request and one comment that may require an update.

@lhutton1
Address review comments
38ef363 
Cleanups in the extension definition.

Change-Id: I48f0be28e8298021a37280c9065fe97b5741fff6
psunn
psunn approved these changes Oct 24, 2025
View reviewed changes
Copy link
Contributor

@psunn psunn left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks

@lhutton1 lhutton1 merged commit 86fd3af into llvm:main Oct 24, 2025
10 checks passed
@lhutton1 lhutton1 deleted the int64-support branch October 24, 2025 10:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@psunn psunn psunn approved these changes

@Jerry-Ge Jerry-Ge Awaiting requested review from Jerry-Ge

Assignees

No one assigned

Labels

mlir:tosa mlir

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@lhutton1 @llvmbot @psunn