commit 42968bf7ab9f7b2d24789ac74f0993f86fe7ed51
author Kuangyuan Chen <chky@google.com> Tue Jul 07 09:31:03 2020 -0700
committer Copybara-Service <iree-copybara-bot@google.com> Tue Jul 07 09:32:42 2020 -0700
tree 7773fcd2b2c5c78dea6bf01192b9a8c95bcf8651
parent a28cf28ec378dd3c07881e9673e5fa379e74849b

Change mlir::tf_saved_model::LookupBoundInput() to return mlir::Operation*.

Change mlir::tf_saved_model::LookupBoundInput() to return mlir::Operation* as there will be different kinds of operations for bound inputs. Also added a helper function LookupBoundInputOfType() to get bound inputs of the expected type.

PiperOrigin-RevId: 319999187

integrations/tensorflow/compiler/TFSavedModelLowerGlobalTensors.cpp

diff --git a/integrations/tensorflow/compiler/TFSavedModelLowerGlobalTensors.cpp b/integrations/tensorflow/compiler/TFSavedModelLowerGlobalTensors.cpp
index 6d8271d..03c878d 100644
--- a/integrations/tensorflow/compiler/TFSavedModelLowerGlobalTensors.cpp
+++ b/integrations/tensorflow/compiler/TFSavedModelLowerGlobalTensors.cpp
@@ -93,8 +93,8 @@
     OpBuilder builder(func.getBody());
     SmallVector<Value, 8> typeConversionWorklist;
     for (int i = 0, e = func.getNumArguments(); i < e; i++) {
-      auto globalTensor =
-          tf_saved_model::LookupBoundInput(func, i, symbolTable);
+      auto globalTensor = tf_saved_model::LookupBoundInputOfType<
+          tf_saved_model::GlobalTensorOp>(func, i, symbolTable);
       if (!globalTensor) {
         continue;
       }