[mlir][Linalg] Allow specifiying zero-rank shaped type operands to linalg.indexed_generic ops.

[hanchung]

Summary

Patch D74638 allows linalg.generic ops to use zero-rank shaped type operands,
this also can be applied to linalg.indexed_generic ops.

[mravishankar]

mlir/test/Dialect/Linalg/loops.mlir

394

You could simplify the test by just having the operands as i32 element types

[hanchung]

Simplify the test by using i32 element types.

[hanchung]

mlir/test/Dialect/Linalg/loops.mlir

394

Sounds good, let's use i32 types.

mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp

static void emitScalarImplementation(ArrayRef<Value> allIvs,
     SmallVector<Value, 4> indexedValues(nLoops + nInputs + nOutputs);
 
     for (unsigned i = 0; i < nLoops; ++i) {
       indexedValues[i] = allIvs[i];
     }
 
     // 1.a. Emit std_load from input views.
     for (unsigned i = 0; i < nInputs; ++i) {
+      Value input = indexedGenericOp.getInput(i);
+      if (!input.getType().cast<ShapedType>().getRank()) {
+        indexedValues[nLoops + i] = std_load(input);
+      } else {
         ValueHandleArray indexing(makeCanonicalAffineApplies(
             b, loc, indexedGenericOp.getInputIndexingMap(i), allIvs));
-      indexedValues[nLoops + i] =
-          std_load(indexedGenericOp.getInput(i), indexing);
+        indexedValues[nLoops + i] = std_load(input, indexing);
+      }
     }
 
     // 1.b. Emit std_load from output views.
     for (unsigned i = 0; i < nOutputs; ++i) {
       ValueHandleArray indexing(makeCanonicalAffineApplies(
           b, loc, indexedGenericOp.getOutputIndexingMap(i), allIvs));
       indexedValues[nLoops + nInputs + i] =
           std_load(indexedGenericOp.getOutputBuffer(i), indexing);

mlir/test/Dialect/Linalg/loops.mlir

 // CHECK:       %[[ijk_int:.*]] = index_cast %[[ijk]] : index to i32
 // CHECK:       %[[ijk_float:.*]] = sitofp %[[ijk_int]] : i32 to f32
 // CHECK:       %[[result_2:.*]] = addf %[[c]], %[[ijk_float]] : f32
 // CHECK:       store %[[result_1]], %{{.*}}[%[[i]], %[[j]], %[[k]]]
 // CHECK:       store %[[result_2]], %{{.*}}[%[[i]], %[[k]], %[[j]]]
 
 // -----
 
-
 #broadcast_access = [
   affine_map<(i, j) -> (0)>,
-  affine_map<(i,j) -> (i,j)>
+  affine_map<(i, j) -> (i, j)>
 ]
 
 #trait_broadcast = {
   args_in = 1,
   args_out = 1,
   indexing_maps = #broadcast_access,
   iterator_types = ["parallel", "parallel"],
   library_call = "some_broadcast_external_fn"
 }
 
-func @generic_op_zero_rank(%arg0 : memref<f32>, %arg1:  memref<3x4xf32>)
+func @generic_op_zero_rank(%arg0: memref<f32>, %arg1: memref<3x4xf32>)
 {
   linalg.generic #trait_broadcast %arg0, %arg1 {
-    ^bb(%a: f32, %b : f32) :
+    ^bb(%a: f32, %b: f32) :
       linalg.yield %a : f32
   } : memref<f32>, memref<3x4xf32>
   return
 }
 
 // CHECK-LABEL: @generic_op_zero_rank
 // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<f32>
 // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<3x4xf32>
 // CHECK: loop.for %[[i:.*]] = {{.*}}
 // CHECK:   loop.for %[[j:.*]] = {{.*}}
 // CHECK:     %[[a:.*]] = load %[[ARG0]][]
 // CHECK:     store %[[a]], %[[ARG1]][%[[i]], %[[j]]]
+
+func @indexed_generic_op_zero_rank(%arg0: memref<i32>, %arg1: memref<3x4xi32>)
+{
+  linalg.indexed_generic #trait_broadcast %arg0, %arg1 {
+    ^bb(%i: index, %j: index, %a: i32, %b: i32) :
+      %ij = addi %i, %j : index
+      %ij_int = index_cast %ij : index to i32
+      %result = addi %a, %ij_int : i32
+      linalg.yield %result : i32
+  } : memref<i32>, memref<3x4xi32>
+  return
+}
+
+// CHECK-LABEL: @indexed_generic_op_zero_rank
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<i32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<3x4xi32>
+// CHECK: loop.for %[[i:.*]] = {{.*}}
+// CHECK:   loop.for %[[j:.*]] = {{.*}}
+// CHECK:     %[[a:.*]] = load %[[ARG0]][
+// CHECK:     %[[ij:.*]] = addi %[[i]], %[[j]] : index
+// CHECK:     %[[ij_int:.*]] = index_cast %[[ij]] : index to i32
+// CHECK:     %[[result:.*]] = addi %[[a]], %[[ij_int]] : i32
+// CHECK:     store %[[result]], %[[ARG1]][%[[i]], %[[j]]]

mlir/test/Dialect/Linalg/roundtrip.mlir

 //  CHECK-SAME:     library_call = "some_external_function_name_1"} %{{.*}}, %{{.*}} {foo = 1 : i64}:
 //  CHECK-SAME:     tensor<?x?xvector<3x4xi4>>, tensor<?x?x?xf32> -> tensor<?x?x?xf32>
 //       CHECK:   return {{.*}} : tensor<?x?x?xf32>
 
 // -----
 
 #broadcast_access = [
   affine_map<(i, j) -> (0)>,
-  affine_map<(i,j) -> (i,j)>
+  affine_map<(i, j) -> (i, j)>
 ]
 
 #trait_broadcast = {
   args_in = 1,
   args_out = 1,
   indexing_maps = #broadcast_access,
   iterator_types = ["parallel", "parallel"],
   library_call = "some_broadcast_external_fn"
 }
 
-func @generic_op_zero_rank(%arg0 : tensor<f32>) ->  (tensor<3x4xf32>)
+func @generic_op_zero_rank(%arg0: tensor<f32>) -> (tensor<3x4xf32>)
 {
   %0 = linalg.generic #trait_broadcast %arg0 {
     ^bb(%a: f32) :
       linalg.yield %a : f32
   } : tensor<f32> -> tensor<3x4xf32>
   return %0 : tensor<3x4xf32>
 }
 
+func @indexed_generic_op_zero_rank(%arg0: tensor<f32>) -> (tensor<3x4xf32>)
+{
+  %0 = linalg.indexed_generic #trait_broadcast %arg0 {
+    ^bb(%i: index, %j: index, %a: f32) :
+      linalg.yield %a : f32
+  } : tensor<f32> -> tensor<3x4xf32>
+  return %0 : tensor<3x4xf32>
+}
+
 // -----
 
 // CHECK-DAG: #[[strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 // CHECK-DAG: #[[strided3D:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>
 
 #accesses = [
   affine_map<(i, j, k) -> (j, i)>,
   affine_map<(i, j, k) -> (i, k, i + j)>