From 221fa96cd43cdcb3af5b03d814f38643874bd794 Mon Sep 17 00:00:00 2001
From: Ahmed Taei <ataei@google.com>
Date: Mon, 30 Mar 2020 17:36:05 -0700
Subject: [PATCH] Fix linalg.generic access of hoisted constants

Summary: Otherwise the added @generic_const_int will fail

Reviewers: nicolasvasilache, rriddle, mravishankar

Subscribers: mehdi_amini, jpienaar, burmako, shauheen, antiagainst, arpith-jacob, mgester, lucyrfox, liufengdb, Joonsoo, grosul1, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D77109
---
 .../Linalg/Transforms/LinalgToLoops.cpp       |  2 +-
 mlir/test/Dialect/Linalg/loops.mlir           | 21 +++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp b/mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp
index 971f1b0ca863..eb2a881308c4 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp
@@ -101,7 +101,7 @@ static void inlineRegionAndEmitStdStore(OpType op,
   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]);
   }
 }
diff --git a/mlir/test/Dialect/Linalg/loops.mlir b/mlir/test/Dialect/Linalg/loops.mlir
index b2cd61791d66..c8d114bee6ae 100644
--- a/mlir/test/Dialect/Linalg/loops.mlir
+++ b/mlir/test/Dialect/Linalg/loops.mlir
@@ -522,3 +522,24 @@ func @indexed_generic_op_1D_reduce(%arg0: memref<?xf32>,
 // 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]]