@@ -241,6 +241,10 @@ limitations under the License.
#ifdef ENABLE_ANNC
#include <annc/kdnn_rewriter.h>
#include <annc/annc_flags.h>
+#include <xnnpack_ops_rewriter.h>
+#include <bisheng-cpu/embedding_simplify.h>
+#include <bisheng-cpu/concat_partition.h>
+#include <bisheng-cpu/reduce_combine.h>
#endif
namespace {
@@ -705,6 +709,14 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn(
pipeline.AddPass<BatchDotSimplification>();
pipeline.AddPass<DotDecomposer>();
+#ifdef ENABLE_ANNC
+ // Rewrite to custom calls with XNNPACK targets.
+ bool enable_xnnpack =
+ xla::GetDebugOptionsFromFlags().xla_cpu_enable_xnnpack();
+ if (enable_xnnpack)
+ pipeline.AddPass<XnnPackOpsRewriter>();
+#endif
+
// Rewrite to custom calls with target as oneDNN library calls.
#if defined(INTEL_MKL) && defined(ENABLE_ONEDNN_V3)
// AOT compiled code runs in single thread.
@@ -865,6 +877,13 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn(
#ifdef ENABLE_ANNC
+ if (getenv("ENABLE_BISHENG_GRAPH_OPT") != NULL) {
+ pipeline.AddPass<ConcatPartition>();
+ pipeline.AddPass<EmbeddingSimplify>();
+ pipeline.AddPass<ReduceCombine>();
+ pipeline.AddPass<HloDCE>();
+ }
+
auto& flags = annc::get_annc_flags();
if (flags.is_enabled("annc-pass")) {
VLOG(1) << "#### Run KDnnFusion Pass" ;
@@ -904,7 +923,14 @@ Status CpuCompiler::RunHloPassesAfterLayoutAssn(
pipeline.AddPass<ReshapeDecomposer>();
// Add a fusion pass now that layout assignment is done.
+#ifdef ENABLE_ANNC
+ if (getenv("ENABLE_BISHENG_GRAPH_OPT") != NULL)
+ pipeline.AddPass<CpuInstructionFusion>(/*may_duplicate=*/false);
+ else
+ pipeline.AddPass<CpuInstructionFusion>(/*may_duplicate=*/true);
+#else
pipeline.AddPass<CpuInstructionFusion>();
+#endif
// The LayoutAssignment pass may leave behind kCopy instructions which are
// duplicate or NOPs, so remove them with algebraic simplification and CSE.