diff --git a/WORKSPACE b/WORKSPACE
index fa8e277..a61ef35 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -15,9 +15,9 @@
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
-LLVM_COMMIT = "6381664580080f015bc0c2ec647853f697cf744a"
+LLVM_COMMIT = "0316f3e64972c919d8bfa2d15b9a4be858530f85"
 
-LLVM_SHA256 = "f16c1757a31dbf9370da3de766e20b238212b5a0e08b897eeda893899ebf70bf"
+LLVM_SHA256 = "ddf0daa23870922c468343c08670e8007358b463cf4ed2c737a94086cc433527"
 
 LLVM_BAZEL_TAG = "llvm-project-{commit}".format(commit = LLVM_COMMIT)
 
diff --git a/build_tools/llvm_version.txt b/build_tools/llvm_version.txt
index 0cea335..70dc8a4 100644
--- a/build_tools/llvm_version.txt
+++ b/build_tools/llvm_version.txt
@@ -1,2 +1,2 @@
-6381664580080f015bc0c2ec647853f697cf744a
+0316f3e64972c919d8bfa2d15b9a4be858530f85
 
diff --git a/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td b/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td
index 2116556..477cfda 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td
+++ b/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td
@@ -438,7 +438,7 @@ class HLOClient_UnaryElementwiseOp<string mnemonic, list<OpTrait> traits,
         SmallVectorImpl<ShapedTypeComponents>& inferedReturnShapes) {
       return failure();
     }
-    LogicalResult reifyReturnTypeShapes(OpBuilder& builder,
+    LogicalResult reifyReturnTypeShapes(OpBuilder& builder, ValueRange operands,
         SmallVectorImpl<Value>& reifiedReturnShapes) {
       return ::mlir::mhlo::deriveShapeFromFirstOperand(&builder, getOperation(),
           &reifiedReturnShapes);
diff --git a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
index b6c4340..6c2949e 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
+++ b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
@@ -139,7 +139,8 @@ class HLO_UnaryElementwiseOp<string mnemonic, list<OpTrait> traits,
       return failure();
     }
     LogicalResult reifyReturnTypeShapes(
-        OpBuilder& builder, SmallVectorImpl<Value>& reifiedReturnShapes) {
+        OpBuilder& builder, ValueRange operands,
+        SmallVectorImpl<Value>& reifiedReturnShapes) {
       return ::mlir::mhlo::deriveShapeFromFirstOperand(&builder, getOperation(),
                                                        &reifiedReturnShapes);
     }
@@ -452,7 +453,8 @@ class HLO_BinaryElementwiseOp<string mnemonic, list<OpTrait> traits> :
       return failure();
     }
     LogicalResult reifyReturnTypeShapes(
-        OpBuilder& builder, SmallVectorImpl<Value>& reifiedReturnShapes) {
+        OpBuilder& builder, ValueRange operands,
+        SmallVectorImpl<Value>& reifiedReturnShapes) {
       return ::mlir::mhlo::deriveShapeFromFirstOperand(&builder, getOperation(),
                                                        &reifiedReturnShapes);
     }
diff --git a/lib/Dialect/mhlo/IR/chlo_ops.cc b/lib/Dialect/mhlo/IR/chlo_ops.cc
index ad149e2..64ee93c 100644
--- a/lib/Dialect/mhlo/IR/chlo_ops.cc
+++ b/lib/Dialect/mhlo/IR/chlo_ops.cc
@@ -204,7 +204,8 @@ LogicalResult BroadcastComplexOp::inferReturnTypeComponents(
                                                     inferedReturnShapes);
 }
 LogicalResult BroadcastComplexOp::reifyReturnTypeShapes(
-    OpBuilder& builder, SmallVectorImpl<Value>& reifiedReturnShapes) {
+    OpBuilder& builder, ValueRange,
+    SmallVectorImpl<Value>& reifiedReturnShapes) {
   return ReifyBroadcastBinaryOpReturnTypeShapes(builder, getOperation(),
                                                 reifiedReturnShapes);
 }
@@ -235,7 +236,8 @@ LogicalResult BroadcastCompareOp::inferReturnTypeComponents(
 }
 
 LogicalResult BroadcastCompareOp::reifyReturnTypeShapes(
-    OpBuilder& builder, SmallVectorImpl<Value>& reifiedReturnShapes) {
+    OpBuilder& builder, ValueRange,
+    SmallVectorImpl<Value>& reifiedReturnShapes) {
   return ReifyBroadcastBinaryOpReturnTypeShapes(builder, getOperation(),
                                                 reifiedReturnShapes);
 }
@@ -293,7 +295,8 @@ LogicalResult IsPosInfOp::inferReturnTypes(
         inferedReturnShapes);                                                 \
   }                                                                           \
   LogicalResult Op::reifyReturnTypeShapes(                                    \
-      OpBuilder& builder, SmallVectorImpl<Value>& reifiedReturnShapes) {      \
+      OpBuilder& builder, ValueRange,                                         \
+      SmallVectorImpl<Value>& reifiedReturnShapes) {                          \
     return ReifyBroadcastBinaryOpReturnTypeShapes(builder, getOperation(),    \
                                                   reifiedReturnShapes);       \
   }
diff --git a/lib/Dialect/mhlo/IR/hlo_ops.cc b/lib/Dialect/mhlo/IR/hlo_ops.cc
index d6855d6..0373b87 100644
--- a/lib/Dialect/mhlo/IR/hlo_ops.cc
+++ b/lib/Dialect/mhlo/IR/hlo_ops.cc
@@ -1002,7 +1002,7 @@ LogicalResult DynamicBroadcastInDimOp::inferReturnTypeComponents(
 }
 
 LogicalResult DynamicBroadcastInDimOp::reifyReturnTypeShapes(
-    OpBuilder&, SmallVectorImpl<Value>& reifiedReturnShapes) {
+    OpBuilder&, ValueRange, SmallVectorImpl<Value>& reifiedReturnShapes) {
   reifiedReturnShapes.push_back(output_dimensions());
   return success();
 }
@@ -2100,7 +2100,8 @@ LogicalResult SelectOp::inferReturnTypeComponents(
 }
 
 LogicalResult SelectOp::reifyReturnTypeShapes(
-    OpBuilder& builder, SmallVectorImpl<Value>& reifiedReturnShapes) {
+    OpBuilder& builder, ValueRange operands,
+    SmallVectorImpl<Value>& reifiedReturnShapes) {
   return deriveShapeFromFirstOperand(&builder, getOperation(),
                                      &reifiedReturnShapes);
 }
@@ -3183,7 +3184,8 @@ LogicalResult CompareOp::inferReturnTypeComponents(
 }
 
 LogicalResult CompareOp::reifyReturnTypeShapes(
-    OpBuilder& builder, SmallVectorImpl<Value>& reifiedReturnShapes) {
+    OpBuilder& builder, ValueRange operands,
+    SmallVectorImpl<Value>& reifiedReturnShapes) {
   return deriveShapeFromFirstOperand(&builder, getOperation(),
                                      &reifiedReturnShapes);
 }
diff --git a/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc b/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc
index 87b2496..e36a821 100644
--- a/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc
+++ b/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc
@@ -107,7 +107,8 @@ LogicalResult ConvertResults(Operation* op, SmallVectorImpl<Value>& results,
     if (!shape_type_op) return failure();
 
     SmallVector<Value, 1> results_shape;
-    auto status = shape_type_op.reifyReturnTypeShapes(rewriter, results_shape);
+    auto status = shape_type_op.reifyReturnTypeShapes(
+        rewriter, shape_type_op->getOperands(), results_shape);
     if (failed(status)) return failure();
     results.push_back(
         InsertDynamicAllocAndDealloc(op->getLoc(), result.value(),
@@ -390,7 +391,8 @@ struct HloToLhloDotGeneralOpConverter
     } else {
       SmallVector<Value, 1> results_shape;
       auto shape_type_op = dyn_cast<InferShapedTypeOpInterface>(op);
-      if (failed(shape_type_op.reifyReturnTypeShapes(rewriter, results_shape)))
+      if (failed(shape_type_op.reifyReturnTypeShapes(
+              rewriter, shape_type_op->getOperands(), results_shape)))
         return failure();
 
       bufferArgs[2] = InsertDynamicAllocAndDealloc(
diff --git a/lib/Dialect/mhlo/transforms/move_up_dynamic_broadcasts_for_fusion.cc b/lib/Dialect/mhlo/transforms/move_up_dynamic_broadcasts_for_fusion.cc
index 82bb69c..51576a7 100644
--- a/lib/Dialect/mhlo/transforms/move_up_dynamic_broadcasts_for_fusion.cc
+++ b/lib/Dialect/mhlo/transforms/move_up_dynamic_broadcasts_for_fusion.cc
@@ -55,7 +55,8 @@ struct ShapeReificationPattern : public OpRewritePattern<shape::ShapeOfOp> {
     if (!shape_origin) return failure();
 
     llvm::SmallVector<Value, 1> reifications;
-    if (failed(shape_origin.reifyReturnTypeShapes(rewriter, reifications)))
+    if (failed(shape_origin.reifyReturnTypeShapes(
+            rewriter, shape_origin->getOperands(), reifications)))
       return failure();
     assert(reifications.size() == 1);
     Value reified_shape = reifications.front();
diff --git a/lib/Dialect/mhlo/transforms/test_infer_shaped_type_pass.cc b/lib/Dialect/mhlo/transforms/test_infer_shaped_type_pass.cc
index d9859eb..d9588b3 100644
--- a/lib/Dialect/mhlo/transforms/test_infer_shaped_type_pass.cc
+++ b/lib/Dialect/mhlo/transforms/test_infer_shaped_type_pass.cc
@@ -74,7 +74,8 @@ struct ReifyReturnTypeShapesPattern : public RewritePattern {
         op->getOperand(0).getDefiningOp());
     if (!defining_op) return failure();
     SmallVector<Value, 4> return_shapes;
-    if (failed(defining_op.reifyReturnTypeShapes(rewriter, return_shapes))) {
+    if (failed(defining_op.reifyReturnTypeShapes(
+            rewriter, defining_op->getOperands(), return_shapes))) {
       return failure();
     }
     rewriter.replaceOp(op, return_shapes);