From 8cc6cddf9f237e352684fa7f48eb542e2a9cb35a Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Sun, 13 Sep 2020 23:51:36 -0700
Subject: [PATCH] Use location instead of unknown loc

Follow up on TODO as placeholder has been added to rewriter.

PiperOrigin-RevId: 331488583
---
 include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.h   | 6 +++---
 include/mlir-hlo/Dialect/mhlo/IR/hlo_utils.td | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.h b/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.h
index 9704f34..282b056 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.h
+++ b/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.h
@@ -48,7 +48,8 @@ class HloClientDialect : public Dialect {
 #include "mlir-hlo/Dialect/mhlo/IR/chlo_ops.h.inc"
 
 template <typename T>
-static Value getConstantLike(OpBuilder& b, T constant, Value val) {
+static Value getConstantLike(OpBuilder& b, Location loc, T constant,
+                             Value val) {
   Type ty = getElementTypeOrSelf(val.getType());
 
   auto getAttr = [&]() -> Attribute {
@@ -56,8 +57,7 @@ static Value getConstantLike(OpBuilder& b, T constant, Value val) {
     if (ty.isa<FloatType>()) return b.getFloatAttr(ty, constant);
     llvm_unreachable("unhandled element type");
   };
-  // TODO(jpienaar): Add ability to pass loc via native call and update.
-  return b.create<ConstantLikeOp>(b.getUnknownLoc(), getAttr(), val);
+  return b.create<ConstantLikeOp>(loc, getAttr(), val);
 }
 
 }  // namespace chlo
diff --git a/include/mlir-hlo/Dialect/mhlo/IR/hlo_utils.td b/include/mlir-hlo/Dialect/mhlo/IR/hlo_utils.td
index c201aef..32940cb 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/hlo_utils.td
+++ b/include/mlir-hlo/Dialect/mhlo/IR/hlo_utils.td
@@ -28,7 +28,7 @@ class ConstantSplat<string value> : NativeCodeCall<
     "hlo::getSplat(&$_builder, $0, " # value # ")">;
 
 class HLO_ConstantLike<string value> : NativeCodeCall<
-    "chlo::getConstantLike($_builder, " # value # ", $0)">;
+    "chlo::getConstantLike($_builder, $_loc, " # value # ", $0)">;
 
 def NullDenseIntElementsAttr : NativeCodeCall<"DenseIntElementsAttr()">;