From 6ed8e02c2bd01fc4fc058faec4ca62e6168ac080 Mon Sep 17 00:00:00 2001
From: Mahmoud Abuzaina <mahmoud.abuzaina@intel.com>
Date: Wed, 4 Oct 2023 20:57:53 -0700
Subject: [PATCH] Updated test to check for fusion

---
 third_party/xla/xla/tests/onednn_layer_norm_test.cc | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/third_party/xla/xla/tests/onednn_layer_norm_test.cc b/third_party/xla/xla/tests/onednn_layer_norm_test.cc
index 5a86619f4b2709..a1cd2224481a85 100644
--- a/third_party/xla/xla/tests/onednn_layer_norm_test.cc
+++ b/third_party/xla/xla/tests/onednn_layer_norm_test.cc
@@ -85,6 +85,15 @@ TEST_F(LayerNormTest, SimpleTest) {
 )";
 
   EXPECT_TRUE(RunAndCompare(layer_norm_module_str, ErrorSpec{1e-4, 1e-4}));
+  MatchOptimizedHlo(layer_norm_module_str,
+                    R"(
+  ; CHECK:     custom_call_target="__onednn$layernorm",
+  ; CHECK:       backend_config={
+  ; CHECK-DAG:     "onednn_layer_norm_config":{
+  ; CHECK-DAG:       "fused_ops":"SCALE_AND_SHIFT"
+  ; CHECK-DAG:   }
+  ; CHECK:     }
+  )");
 }
 
 }  // namespace