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

[mravishankar]

Summary

Fixing a bug where using a zero-rank shaped type operand to
linalg.generic ops hit an unrelated assert. This also meant that
lowering the operation to loops was not supported. Adding roundtrip
tests and lowering to loops test for zero-rank shaped type operand
with fixes to make the test pass.

[mravishankar]

Removing errant ;

[nicolasvasilache]

Corner cases are the worst :)
Thanks a lot Mahesh, very cool!

[rriddle]

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

242

nit: Wrap this in {} to match the else.

mlir/lib/IR/AffineMap.cpp

354

nit: Can you wrap the predicate in ()?

mlir/test/Dialect/Linalg/loops.mlir

392

nit: Add a newline here.

[mravishankar]

Addressing comments and fixing failing test

[mravishankar]

Change to test to check for scalar load

[hanchung]

Just few formatting nits while going through the examples.

mlir/test/Dialect/Linalg/loops.mlir

365

Add a space after comma

376

Remove the space right after %arg0.
And there are two spaces after "%arg1:", remove one.

379

s/%b :/ %b:

mlir/test/Dialect/Linalg/roundtrip.mlir

350

Add a space after comma.

361

Remove the space between "%arg0" and ":".

361

There are two spaces after "->", remove one.

[hanchung]

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

346–351

This part should also be updated with

  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(input, indexing);
  }
}

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

if (m.getNumDims() != nLoops)
              << idx << " to have " << nLoops
              << " dim(s) to match the number of loops";
 
     if (m.getNumResults() == 1 && view.getRank() == 0) {
       auto cst = m.getResult(0).template dyn_cast<AffineConstantExpr>();
       if (!cst || cst.getValue() != 0)
         return op.emitOpError("expected indexing_map #")
                << idx << " to be 0 to match 0-D view: " << view;
-    }
-
-    if (m.getNumResults() != view.getRank())
+    } else if (m.getNumResults() != view.getRank()) {
       return op.emitOpError("expected indexing_map #")
              << idx << " results to match view rank: " << view;
     }
+  }
 
   auto concatMap = concatAffineMaps(indexingMaps);
   auto aggregateMap = inversePermutation(concatMap);
   if (!aggregateMap)
     return op.emitOpError("expected the concatenation of maps in indexing_map "
                           "to be invertible");
 
   return success();

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

static void emitScalarImplementation(ArrayRef<Value> allIvs,
     auto loc = ScopedContext::getLocation();
     using edsc::intrinsics::detail::ValueHandleArray;
     unsigned nInputs = genericOp.getNumInputs();
     unsigned nOutputs = genericOp.getNumOutputs();
     SmallVector<Value, 4> indexedValues(nInputs + nOutputs);
 
     // 1.a. Emit std_load from input views.
     for (unsigned i = 0; i < nInputs; ++i) {
+      Value input = genericOp.getInput(i);
+      if (!input.getType().cast<ShapedType>().getRank()) {
+        indexedValues[i] = std_load(input);
+      } else {
         ValueHandleArray indexing(makeCanonicalAffineApplies(
             b, loc, genericOp.getInputIndexingMap(i), allIvs));
-      indexedValues[i] = std_load(genericOp.getInput(i), indexing);
+        indexedValues[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, genericOp.getOutputIndexingMap(i), allIvs));
       indexedValues[nInputs + i] =
           std_load(genericOp.getOutputBuffer(i), indexing);
static void emitScalarImplementation(ArrayRef<Value> allIvs,
     unsigned nLoops = allIvs.size();
     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) {
       ValueHandleArray indexing(makeCanonicalAffineApplies(
           b, loc, indexedGenericOp.getInputIndexingMap(i), allIvs));
       indexedValues[nLoops + i] =
           std_load(indexedGenericOp.getInput(i), 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/lib/IR/AffineMap.cpp

AffineMap mlir::inversePermutation(AffineMap map) {
   if (seenExprs.size() != map.getNumInputs())
     return AffineMap();
   return AffineMap::get(map.getNumResults(), 0, seenExprs);
 }
 
 AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
   unsigned numResults = 0;
   for (auto m : maps)
-    numResults += m ? m.getNumResults() : 0;
+    numResults += (m && !m.isSingleConstant()) ? m.getNumResults() : 0;
   unsigned numDims = 0;
   SmallVector<AffineExpr, 8> results;
   results.reserve(numResults);
   for (auto m : maps) {
-    if (!m)
+    if (!m || m.isSingleConstant())
       continue;
     assert(m.getNumSymbols() == 0 && "expected map without symbols");
     results.append(m.getResults().begin(), m.getResults().end());
     numDims = std::max(m.getNumDims(), numDims);
   }
   return numDims == 0 ? AffineMap() : AffineMap::get(numDims, 0, results);
 }
 

mlir/test/Dialect/Linalg/loops.mlir

 // CHECK:       %[[result_1:.*]] = mulf %[[a]], %[[b]] : f32
 // CHECK:       %[[ij:.*]] = addi %[[i]], %[[j]] : index
 // CHECK:       %[[ijk:.*]] = addi %[[ij]], %[[k]] : index
 // 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)>
+]
+
+#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>)
+{
+  linalg.generic #trait_broadcast %arg0, %arg1 {
+    ^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]]]

mlir/test/Dialect/Linalg/roundtrip.mlir

 //       CHECK:   linalg.indexed_generic {args_in = 2 : i64, args_out = 1 : i64, fun = @foo,
 //  CHECK-SAME:     indexing_maps = [#{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel"],
 //  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)>
+]
+
+#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>)
+{
+  %0 = linalg.generic #trait_broadcast %arg0 {
+    ^bb(%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)>
 ]