1 files changed, 8 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/cpu/conv_canonicalization.h b/tensorflow/compiler/xla/service/cpu/conv_canonicalization.h
index 9b2c3d82eb..e6fd1499ed 100644
--- a/tensorflow/compiler/xla/service/cpu/conv_canonicalization.h
+++ b/tensorflow/compiler/xla/service/cpu/conv_canonicalization.h
@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CONV_CANONICALIZATION_H_
 #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CONV_CANONICALIZATION_H_
 
+#include "tensorflow/compiler/xla/service/cpu/target_machine_features.h"
 #include "tensorflow/compiler/xla/service/hlo_module.h"
 #include "tensorflow/compiler/xla/service/hlo_pass_interface.h"
 
@@ -32,12 +33,19 @@ namespace cpu {
 // convolutions can run faster.
 class ConvCanonicalization : public HloPassInterface {
  public:
+  explicit ConvCanonicalization(
+      const TargetMachineFeatures* target_machine_features)
+      : target_machine_features_(*target_machine_features) {}
+
   ~ConvCanonicalization() override {}
   tensorflow::StringPiece name() const override {
     return "convolution-canonicalization";
   }
 
   StatusOr<bool> Run(HloModule* module) override;
+
+ private:
+  const TargetMachineFeatures& target_machine_features_;
 };
 
 }  // namespace cpu