[mlir][Linalg] Uniformize linalg.generic with named ops.

This revision allows representing a reduction at the level of linalg on tensors for generic ops by uniformizing with the named ops approach.

1 files changed, 16 insertions, 16 deletions

diff --git a/mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir b/mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir
index d437ab1..0fac017 100644
--- a/mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir
+++ b/mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir
@@ -5,8 +5,6 @@
 //===----------------------------------------------------------------------===//
 
 #single_workgroup_reduction_trait = {
-  args_in = 1,
-  args_out = 1,
   iterator_types = ["reduction"],
   indexing_maps = [
     affine_map<(i) -> (i)>,
@@ -49,11 +47,13 @@ module attributes {
 func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>) attributes {
   spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}
 } {
-  linalg.generic #single_workgroup_reduction_trait %input, %output {
+  linalg.generic #single_workgroup_reduction_trait
+      ins(%input : memref<16xi32>)
+     outs(%output : memref<1xi32>) {
     ^bb(%in: i32, %out: i32):
       %sum = addi %in, %out : i32
       linalg.yield %sum : i32
-  } : memref<16xi32>, memref<1xi32>
+  }
   spv.Return
 }
 }
@@ -63,8 +63,6 @@ func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>)
 // Missing shader entry point ABI
 
 #single_workgroup_reduction_trait = {
-  args_in = 1,
-  args_out = 1,
   iterator_types = ["reduction"],
   indexing_maps = [
     affine_map<(i) -> (i)>,
@@ -78,11 +76,13 @@ module attributes {
 } {
 func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>) {
   // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
-  linalg.generic #single_workgroup_reduction_trait %input, %output {
+  linalg.generic #single_workgroup_reduction_trait
+      ins(%input : memref<16xi32>)
+     outs(%output : memref<1xi32>) {
     ^bb(%in: i32, %out: i32):
       %sum = addi %in, %out : i32
       linalg.yield %sum : i32
-  } : memref<16xi32>, memref<1xi32>
+  }
   return
 }
 }
@@ -92,8 +92,6 @@ func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>)
 // Mismatch between shader entry point ABI and input memref shape
 
 #single_workgroup_reduction_trait = {
-  args_in = 1,
-  args_out = 1,
   iterator_types = ["reduction"],
   indexing_maps = [
     affine_map<(i) -> (i)>,
@@ -109,11 +107,13 @@ func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>)
   spv.entry_point_abi = {local_size = dense<[32, 1, 1]>: vector<3xi32>}
 } {
   // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
-  linalg.generic #single_workgroup_reduction_trait %input, %output {
+  linalg.generic #single_workgroup_reduction_trait
+      ins(%input : memref<16xi32>)
+     outs(%output : memref<1xi32>) {
     ^bb(%in: i32, %out: i32):
       %sum = addi %in, %out : i32
       linalg.yield %sum : i32
-  } : memref<16xi32>, memref<1xi32>
+  }
   spv.Return
 }
 }
@@ -123,8 +123,6 @@ func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>)
 // Unsupported multi-dimension input memref
 
 #single_workgroup_reduction_trait = {
-  args_in = 1,
-  args_out = 1,
   iterator_types = ["parallel", "reduction"],
   indexing_maps = [
     affine_map<(i, j) -> (i, j)>,
@@ -140,11 +138,13 @@ func @single_workgroup_reduction(%input: memref<16x8xi32>, %output: memref<16xi3
   spv.entry_point_abi = {local_size = dense<[16, 8, 1]>: vector<3xi32>}
 } {
   // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
-  linalg.generic #single_workgroup_reduction_trait %input, %output {
+  linalg.generic #single_workgroup_reduction_trait
+      ins(%input : memref<16x8xi32>)
+     outs(%output : memref<16xi32>) {
     ^bb(%in: i32, %out: i32):
       %sum = addi %in, %out : i32
       linalg.yield %sum : i32
-  } : memref<16x8xi32>, memref<16xi32>
+  }
   spv.Return
 }
 }