Fix linalg.generic access of hoisted constants

[asaadaldien]

Summary

Otherwise the added @generic_const_int will fail

[nicolasvasilache]

Thanks for the fix!

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

 //===- LinalgToLoops.cpp - conversion from Linalg library ops to loops-----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
for (auto &op : block.without_terminator()) {
     auto *newOp = b.clone(op, map);
     map.map(op.getResults(), newOp->getResults());
   }
 
   Operation &terminator = block.back();
   assert(isa<YieldOp>(terminator) &&
          "expected an yield op in the end of the region");
   for (unsigned i = 0, e = terminator.getNumOperands(); i < e; ++i) {
-    std_store(map.lookup(terminator.getOperand(i)), outputBuffers[i],
+    std_store(map.lookupOrDefault(terminator.getOperand(i)), outputBuffers[i],
               indexing[i]);
   }
 }
 
 namespace {
 template <typename IndexedValueType, typename LinalgOpType>
 class LinalgScopedEmitter {};
 

mlir/test/Dialect/Linalg/loops.mlir

 // 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]][]
+
+#trait_const_fill = {
+  args_in = 0,
+  args_out = 1,
+  indexing_maps = [affine_map<(i) -> (i)>],
+  iterator_types = ["parallel"],
+  library_call = "some_external_fn"
+}
+func @generic_const_init(%arg0: memref<?xf32>) {
+  %cst = constant 1.0 : f32
+  linalg.generic #trait_const_fill %arg0 {
+    ^bb0(%arg1: f32):   // no predecessors
+      linalg.yield %cst : f32
+    }: memref<?xf32>
+    return
+}
+// CHECK-LABEL: @generic_const_init
+// CHECK-SAME: %[[ARG0:.*]]: memref<?xf32>
+// CHECK: %[[CONST:.*]] = constant 1.000000e+00 : f32
+// CHECK: loop.for %[[i:.*]] = {{.*}}
+// CHECK:   store %[[CONST]], %[[ARG0]]