[Mlir-commits] [mlir] [mlir][tosa] Allow integer gather/scatter ops in fp profile (PR #183342)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-tosa

Author: Luke Hutton (lhutton1)

<details>
<summary>Changes</summary>

This commit updates profile compliance to allow integer gather and scatter operations to be used with the floating point profile. This update aligns with the specification change: https://github.com/arm/tosa-specification/pull/35.

---
Full diff: https://github.com/llvm/llvm-project/pull/183342.diff


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc (+16-6) 
- (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir (+96) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
index 3bdb3805dc6a2..7ea0d134941c7 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
@@ -370,10 +370,16 @@ profileComplianceMap = {
         {{i16T, i32T, i16T}, SpecificationVersion::V_1_0},
         {{i32T, i32T, i32T}, SpecificationVersion::V_1_0}}},
       {{Profile::pro_fp},
-       {{{fp16T, i32T, fp16T}, SpecificationVersion::V_1_0},
+       {{{i8T, i32T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i16T, i32T, i16T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i32T, i32T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp16T, i32T, fp16T}, SpecificationVersion::V_1_0},
         {{fp32T, i32T, fp32T}, SpecificationVersion::V_1_0}}},
       {{Profile::pro_fp, Profile::pro_int},
-       {{{boolT, i32T, boolT}, SpecificationVersion::V_1_1_DRAFT}},
+       {{{boolT, i32T, boolT}, SpecificationVersion::V_1_1_DRAFT},
+        {{i8T, i64T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i16T, i64T, i16T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i32T, i64T, i32T}, SpecificationVersion::V_1_1_DRAFT}},
        anyOf}}},
     {"tosa.scatter",
      {{{Profile::pro_int},
@@ -381,12 +387,16 @@ profileComplianceMap = {
         {{i16T, i32T, i16T, i16T}, SpecificationVersion::V_1_0},
         {{i32T, i32T, i32T, i32T}, SpecificationVersion::V_1_0}}},
       {{Profile::pro_fp},
-       {{{fp16T, i32T, fp16T, fp16T}, SpecificationVersion::V_1_0},
+       {{{i8T, i32T, i8T, i8T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i16T, i32T, i16T, i16T}, SpecificationVersion::V_1_1_DRAFT},
+        {{i32T, i32T, i32T, i32T}, SpecificationVersion::V_1_1_DRAFT},
+        {{fp16T, i32T, fp16T, fp16T}, SpecificationVersion::V_1_0},
         {{fp32T, i32T, fp32T, fp32T}, SpecificationVersion::V_1_0}}},
       {{Profile::pro_fp, Profile::pro_int},
-       {
-           {{boolT, i32T, boolT, boolT}, SpecificationVersion::V_1_1_DRAFT},
-       },
+       {{{boolT, i32T, boolT, boolT}, SpecificationVersion::V_1_1_DRAFT},
+        {{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}},
        anyOf}}},
     {"tosa.resize",
      {{{Profile::pro_int},
diff --git a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
index fa03abcccacd8..4a6ab456744db 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -215,6 +215,102 @@ func.func @test_scatter_bool_i32(%arg0: tensor<13x52x3xi1>, %arg1: tensor<13x26x
 
 // -----
 
+// CHECK-LABEL: test_gather_i8_i32_indices
+func.func @test_gather_i8_i32_indices(%arg0: tensor<13x21x3xi8>, %arg1: tensor<13x26xi32>) -> tensor<13x26x3xi8> {
+  %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x3xi8>, tensor<13x26xi32>) -> tensor<13x26x3xi8>
+  return %0 : tensor<13x26x3xi8>
+}
+
+// -----
+
+// CHECK-LABEL: test_scatter_i8_i32_indices
+func.func @test_scatter_i8_i32_indices(%arg0: tensor<13x27x3xi8>, %arg1: tensor<13x26xi32>, %arg2: tensor<13x26x3xi8>) -> tensor<13x27x3xi8> {
+  %0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x27x3xi8>, tensor<13x26xi32>, tensor<13x26x3xi8>) -> tensor<13x27x3xi8>
+  return %0 : tensor<13x27x3xi8>
+}
+
+// -----
+
+// CHECK-LABEL: test_gather_i16_i32_indices
+func.func @test_gather_i16_i32_indices(%arg0: tensor<13x21x3xi16>, %arg1: tensor<13x26xi32>) -> tensor<13x26x3xi16> {
+  %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x3xi16>, tensor<13x26xi32>) -> tensor<13x26x3xi16>
+  return %0 : tensor<13x26x3xi16>
+}
+
+// -----
+
+// CHECK-LABEL: test_scatter_i16_i32_indices
+func.func @test_scatter_i16_i32_indices(%arg0: tensor<13x27x3xi16>, %arg1: tensor<13x26xi32>, %arg2: tensor<13x26x3xi16>) -> tensor<13x27x3xi16> {
+  %0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x27x3xi16>, tensor<13x26xi32>, tensor<13x26x3xi16>) -> tensor<13x27x3xi16>
+  return %0 : tensor<13x27x3xi16>
+}
+
+// -----
+
+// CHECK-LABEL: test_gather_i32_i32_indices
+func.func @test_gather_i32_i32_indices(%arg0: tensor<13x21x3xi32>, %arg1: tensor<13x26xi32>) -> tensor<13x26x3xi32> {
+  %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x3xi32>, tensor<13x26xi32>) -> tensor<13x26x3xi32>
+  return %0 : tensor<13x26x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: test_scatter_i32_i32_indices
+func.func @test_scatter_i32_i32_indices(%arg0: tensor<13x27x3xi32>, %arg1: tensor<13x26xi32>, %arg2: tensor<13x26x3xi32>) -> tensor<13x27x3xi32> {
+  %0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x27x3xi32>, tensor<13x26xi32>, tensor<13x26x3xi32>) -> tensor<13x27x3xi32>
+  return %0 : tensor<13x27x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: test_gather_i8_i64_indices
+func.func @test_gather_i8_i64_indices(%arg0: tensor<13x21x3xi8>, %arg1: tensor<13x26xi64>) -> tensor<13x26x3xi8> {
+  %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x3xi8>, tensor<13x26xi64>) -> tensor<13x26x3xi8>
+  return %0 : tensor<13x26x3xi8>
+}
+
+// -----
+
+// CHECK-LABEL: test_scatter_i8_i64_indices
+func.func @test_scatter_i8_i64_indices(%arg0: tensor<13x27x3xi8>, %arg1: tensor<13x26xi64>, %arg2: tensor<13x26x3xi8>) -> tensor<13x27x3xi8> {
+  %0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x27x3xi8>, tensor<13x26xi64>, tensor<13x26x3xi8>) -> tensor<13x27x3xi8>
+  return %0 : tensor<13x27x3xi8>
+}
+
+// -----
+
+// CHECK-LABEL: test_gather_i16_i64_indices
+func.func @test_gather_i16_i64_indices(%arg0: tensor<13x21x3xi16>, %arg1: tensor<13x26xi64>) -> tensor<13x26x3xi16> {
+  %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x3xi16>, tensor<13x26xi64>) -> tensor<13x26x3xi16>
+  return %0 : tensor<13x26x3xi16>
+}
+
+// -----
+
+// CHECK-LABEL: test_scatter_i16_i64_indices
+func.func @test_scatter_i16_i64_indices(%arg0: tensor<13x27x3xi16>, %arg1: tensor<13x26xi64>, %arg2: tensor<13x26x3xi16>) -> tensor<13x27x3xi16> {
+  %0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x27x3xi16>, tensor<13x26xi64>, tensor<13x26x3xi16>) -> tensor<13x27x3xi16>
+  return %0 : tensor<13x27x3xi16>
+}
+
+// -----
+
+// CHECK-LABEL: test_gather_i32_i64_indices
+func.func @test_gather_i32_i64_indices(%arg0: tensor<13x21x3xi32>, %arg1: tensor<13x26xi64>) -> tensor<13x26x3xi32> {
+  %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x3xi32>, tensor<13x26xi64>) -> tensor<13x26x3xi32>
+  return %0 : tensor<13x26x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: test_scatter_i32_i64_indices
+func.func @test_scatter_i32_i64_indices(%arg0: tensor<13x27x3xi32>, %arg1: tensor<13x26xi64>, %arg2: tensor<13x26x3xi32>) -> tensor<13x27x3xi32> {
+  %0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x27x3xi32>, tensor<13x26xi64>, tensor<13x26x3xi32>) -> tensor<13x27x3xi32>
+  return %0 : tensor<13x27x3xi32>
+}
+
+// -----
+
 // CHECK-LABEL: test_cast_bool_fp32
 func.func @test_cast_bool_fp32(%arg0: tensor<13x21x3xi1>) -> tensor<13x21x3xf32> {
   %0 = tosa.cast %arg0 : (tensor<13x21x3xi1>) -> tensor<13x21x3xf32>

``````````

</details>


https://github.com/llvm/llvm-project/pull/183342