[mlir][Linalg] Fix load/store operations generated while lower loops when output has zero rank.

[mravishankar]

Summary

While lowering to loops, no indices should be used in the load/store
operation if the buffer is zero-rank.

[nicolasvasilache]

Let's push this forward to unblock progress, in the future the understanding is that linalg ops will not see 0-D memrefs anymore.

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

257

Can we invert the conditional and make it if (X > 0) {} else {} ?
I find it more readable than the not test.

[mravishankar]

Updating based on comments

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

for (unsigned i = 0; i < nInputs; ++i) {
         ValueHandleArray indexing(makeCanonicalAffineApplies(
             b, loc, genericOp.getInputIndexingMap(i), allIvs));
         indexedValues[i] = std_load(input, indexing);
       }
     }
 
     // 1.b. Emit std_load from output views.
     for (unsigned i = 0; i < nOutputs; ++i) {
+      Value output = genericOp.getOutputBuffer(i);
+      if (!output.getType().cast<ShapedType>().getRank()) {
+        indexedValues[nInputs + i] = std_load(output);
+      } else {
         ValueHandleArray indexing(makeCanonicalAffineApplies(
             b, loc, genericOp.getOutputIndexingMap(i), allIvs));
-      indexedValues[nInputs + i] =
-          std_load(genericOp.getOutputBuffer(i), indexing);
+        indexedValues[nInputs + i] = std_load(output, indexing);
+      }
     }
 
     auto funcOp = genericOp.getFunction();
     if (funcOp) {
       // 2. Emit call.
       Operation *callOp = std_call(funcOp, indexedValues);
       assert(callOp->getNumResults() == genericOp.getNumOutputs());
 
       // 3. Emit std_store.
       for (unsigned i = 0; i < nOutputs; ++i) {
+        Value output = genericOp.getOutputBuffer(i);
+        if (!output.getType().cast<ShapedType>().getRank()) {
+          std_store(callOp->getResult(i), output);
+        } else {
           ValueHandleArray indexing(makeCanonicalAffineApplies(
               b, loc, genericOp.getOutputIndexingMap(i), allIvs));
-        std_store(callOp->getResult(i), genericOp.getOutputBuffer(i), indexing);
+          std_store(callOp->getResult(i), output, indexing);
+        }
       }
       return;
     }
     // TODO(ntv): When a region inliner exists, use it.
     // 2. Inline region, currently only works for a single basic block.
     BlockAndValueMapping map;
     auto &block = genericOp.region().front();
     map.map(block.getArguments(), indexedValues);
     for (auto &op : block.without_terminator()) {
       assert(op.getNumRegions() == 0);
       auto *newOp = b.clone(op, map);
       map.map(op.getResults(), newOp->getResults());
     }
 
     // 3. Emit std_store.
     auto *yieldOp = cast<YieldOp>(block.back()).getOperation();
     assert(yieldOp->getNumOperands() == nOutputs);
     for (unsigned i = 0; i < nOutputs; ++i) {
+      Value output = genericOp.getOutputBuffer(i);
+      if (!output.getType().cast<ShapedType>().getRank()) {
+        std_store(map.lookup(yieldOp->getOperand(i)), output);
+      } else {
         ValueHandleArray indexing(makeCanonicalAffineApplies(
             b, loc, genericOp.getOutputIndexingMap(i), allIvs));
         std_store(map.lookup(yieldOp->getOperand(i)),
                   genericOp.getOutputBuffer(i), indexing);
       }
     }
+  }
 };
 
 // Emits the MLIR for the scalar part of the indexed generic op by:
 //   1. Emitting std_load and std_store ops for each input and output view in
 //      order. This is achieved by applying the appropriate input or output map
 //      to the enclosing induction variables.
 //   2. Emitting a call to `op.fun()` that takes as arguments the induction
 //      variables and the scalars from point 1. above.
for (unsigned i = 0; i < nInputs; ++i) {
         ValueHandleArray indexing(makeCanonicalAffineApplies(
             b, loc, indexedGenericOp.getInputIndexingMap(i), allIvs));
         indexedValues[nLoops + i] = std_load(input, indexing);
       }
     }
 
     // 1.b. Emit std_load from output views.
     for (unsigned i = 0; i < nOutputs; ++i) {
+      Value output = indexedGenericOp.getOutputBuffer(i);
+      if (!output.getType().cast<ShapedType>().getRank()) {
+        indexedValues[nLoops + nInputs + i] = std_load(output);
+      } else {
         ValueHandleArray indexing(makeCanonicalAffineApplies(
             b, loc, indexedGenericOp.getOutputIndexingMap(i), allIvs));
-      indexedValues[nLoops + nInputs + i] =
-          std_load(indexedGenericOp.getOutputBuffer(i), indexing);
+        indexedValues[nLoops + nInputs + i] = std_load(output, indexing);
+      }
     }
 
     if (auto funcOp = indexedGenericOp.getFunction()) {
       // 2. Emit call.
       Operation *callOp = std_call(funcOp, indexedValues);
       assert(callOp->getNumResults() == indexedGenericOp.getNumOutputs());
 
       // 3. Emit std_store.
       for (unsigned i = 0; i < nOutputs; ++i) {
+        Value output = indexedGenericOp.getOutputBuffer(i);
+        if (!output.getType().cast<ShapedType>().getRank()) {
+          std_store(callOp->getResult(i), output);
+        } else {
           ValueHandleArray indexing(makeCanonicalAffineApplies(
               b, loc, indexedGenericOp.getOutputIndexingMap(i), allIvs));
-        std_store(callOp->getResult(i), indexedGenericOp.getOutputBuffer(i),
-                  indexing);
+          std_store(callOp->getResult(i), output, indexing);
+        }
       }
       return;
     }
     // TODO(ntv): When a region inliner exists, use it.
     // 2. Inline region, currently only works for a single basic block.
     BlockAndValueMapping map;
     auto &block = indexedGenericOp.region().front();
     map.map(block.getArguments(), indexedValues);
     for (auto &op : block.without_terminator()) {
       assert(op.getNumRegions() == 0);
       auto *newOp = b.clone(op, map);
       map.map(op.getResults(), newOp->getResults());
     }
 
     // 3. Emit std_store.
     auto *yieldOp = cast<YieldOp>(block.back()).getOperation();
     assert(yieldOp->getNumOperands() == nOutputs);
     for (unsigned i = 0; i < nOutputs; ++i) {
+      Value output = indexedGenericOp.getOutputBuffer(i);
+      if (!output.getType().cast<ShapedType>().getRank()) {
+        std_store(map.lookup(yieldOp->getOperand(i)), output);
+      } else {
         ValueHandleArray indexing(makeCanonicalAffineApplies(
             b, loc, indexedGenericOp.getOutputIndexingMap(i), allIvs));
-      std_store(map.lookup(yieldOp->getOperand(i)),
-                indexedGenericOp.getOutputBuffer(i), indexing);
+        std_store(map.lookup(yieldOp->getOperand(i)), output, indexing);
+      }
     }
   }
 };
 
 // This struct is for factoring out the implementation and support template
 // instantiations in the following 2 cases:
 //   1. Appending to a list of patterns via RewritePatternList.
 //   2. Direct invocation via `linalgOpToLoops` and `linalgOpToAffineLoops`.

mlir/test/Dialect/Linalg/loops.mlir

 // 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]]]
+
+#reduce_1D_access = [
+  affine_map<(i) -> (i)>,
+  affine_map<(i) -> (0)>
+]
+
+#trait_reduce_1D = {
+  args_in = 1,
+  args_out = 1,
+  indexing_maps = #reduce_1D_access,
+  iterator_types = ["reduction"],
+  library_call = "some_reduce_external_fn"
+}
+
+func @generic_op_1D_reduce(%arg0: memref<?xf32>, %arg1: memref<f32>)
+{
+  linalg.generic #trait_reduce_1D %arg0, %arg1 {
+    ^bb(%a: f32, %b: f32) :
+      %0 = addf %a, %b : f32
+      linalg.yield %0 : f32
+  } : memref<?xf32>, memref<f32>
+  return
+}
+// CHECK-LABEL: @generic_op_1D_reduce
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<f32>
+// CHECK: loop.for %[[i:.*]] = {{.*}}
+// CHECK:   %[[a:.*]] = load %[[ARG0]][%[[i]]]
+// CHECK:   %[[b:.*]] = load %[[ARG1]][]
+// CHECK:   %[[c:.*]] = addf %[[a]], %[[b]] : f32
+// CHECK:   store %[[c]], %[[ARG1]][]
+
+
+#reduce_init_1D_access = [
+  affine_map<(i) -> (i)>,
+  affine_map<(i) -> (0)>,
+  affine_map<(i) -> (0)>
+]
+
+#trait_reduce_init_1D = {
+  args_in = 2,
+  args_out = 1,
+  indexing_maps = #reduce_init_1D_access,
+  iterator_types = ["reduction"],
+  library_call = "some_reduce_external_fn"
+}
+
+func @indexed_generic_op_1D_reduce(%arg0: memref<?xf32>,
+                                   %arg1: memref<f32>,
+                                   %arg2: memref<f32>)
+{
+  linalg.indexed_generic #trait_reduce_init_1D %arg0, %arg1, %arg2 {
+    ^bb(%i : index, %a: f32, %b: f32, %c: f32) :
+      %0 = constant 0 : index
+      %1 = cmpi "eq", %0, %i : index
+      %2 = select %1, %b, %c : f32
+      %3 = addf %a, %2 : f32
+      linalg.yield %3 : f32
+  } : memref<?xf32>, memref<f32>, memref<f32>
+  return
+}
+// CHECK-LABEL: @indexed_generic_op_1D_reduce
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<f32>
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: memref<f32>
+// CHECK: loop.for %[[i:.*]] = {{.*}}
+// CHECK:   %[[a:.*]] = load %[[ARG0]][%[[i]]]
+// CHECK:   %[[b:.*]] = load %[[ARG1]][]
+// CHECK:   %[[c:.*]] = load %[[ARG2]][]
+// CHECK:   %[[d:.*]] = select %{{.*}}, %[[b]], %[[c]]
+// CHECK:   %[[e:.*]] = addf %[[a]], %[[d]]
+// CHECK:   store %[[e]], %[[ARG2]][]