From 0cf515de79041c58c73397b2c31dc3065a0f8ea0 Mon Sep 17 00:00:00 2001
From: Amir Bishara <amir.bishara@mobileye.com>
Date: Wed, 29 Nov 2023 21:13:54 +0200
Subject: [PATCH] [mlir][linalg] Edit the yieldOutputs method's builder

---
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp                      | 4 ++--
 .../Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp    | 1 +
 mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp   | 2 +-
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 58af9995548e9..11ebbabe97737 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -480,8 +480,8 @@ class RegionBuilderHelper {
     llvm_unreachable("unsupported type conversion function");
   }
 
-  void yieldOutputs(ValueRange values) {
-    OpBuilder builder = getBuilder();
+  void yieldOutputs(OpBuilder builder, ValueRange values) {
+    builder.setInsertionPointToEnd(&block);
     Location loc = builder.getUnknownLoc();
     builder.create<YieldOp>(loc, values);
   }
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
index 375e10f9068e4..e14efd8def092 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
@@ -260,6 +260,7 @@ struct SparseTensorCodegenPass
     // The following operations and dialects may be introduced by the
     // codegen rules, and are therefore marked as legal.
     target.addLegalOp<linalg::FillOp>();
+    target.addLegalOp<linalg::YieldOp>();
     target.addLegalDialect<
         arith::ArithDialect, bufferization::BufferizationDialect,
         complex::ComplexDialect, memref::MemRefDialect, scf::SCFDialect>();
diff --git a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
index fb3c9d48f9a98..cd8398b973a56 100644
--- a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
+++ b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
@@ -1012,7 +1012,7 @@ void {0}::regionBuilder(ImplicitLocOpBuilder &b,
   SmallVector<Value> yields;
   {2}
   {3}
-  helper.yieldOutputs(yields);
+  helper.yieldOutputs(b, yields);
 }
 )FMT";
     auto &args = opConfig.structuredOp->args;