From 699f6c62ab00a95bca8960f10bf1390cd2c5a8ce Mon Sep 17 00:00:00 2001 From: Ilia Sergachev Date: Tue, 14 May 2024 02:59:44 -0700 Subject: [PATCH] PR #12434: [GPU] Fix OSS compilation problems in a previously disabled test. Imported from GitHub PR https://github.com/openxla/xla/pull/12434 Copybara import of the project: -- 723c9bb29adfcc33c015b74f90ce8024c2f79255 by Ilia Sergachev : [GPU] Fix OSS compilation problems in a previously disabled test. Merging this change closes #12434 FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/12434 from openxla:fix_triton_test 723c9bb29adfcc33c015b74f90ce8024c2f79255 PiperOrigin-RevId: 633509548 --- .../xla/xla/service/gpu/ir_emitter_triton_test.cc | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc index 8ebf9e94386885..a04bccac2f73ee 100644 --- a/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc +++ b/third_party/xla/xla/service/gpu/ir_emitter_triton_test.cc @@ -284,7 +284,7 @@ ENTRY e { })"; TritonGemmConfig config(16, 16, 32, 1, 1, 1); - EXPECT_OK( + TF_EXPECT_OK( CreateTritonIrAndFileCheck(kHloText, config, EmitMatMul, "triton_dot", R"( CHECK: tt.func @triton_fn(%[[LHS:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[RHS:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[OUT:.*]]: !tt.ptr {tt.divisibility = 16 : i32}) { CHECK-DAG: %[[ZERO_KN:.*]] = arith.constant dense<0.000000e+00> : tensor<32x16xf32> @@ -1083,8 +1083,8 @@ ENTRY main { ParseAndReturnVerifiedModule(kHloText)); TritonGemmConfig config(16, 64, 32, 1, 1, 1); - ASSERT_OK(CreateTritonIrAndFileCheck(kHloText, config, EmitSoftMax, - "triton_softmax_computation", R"( + TF_ASSERT_OK(CreateTritonIrAndFileCheck(kHloText, config, EmitSoftMax, + "triton_softmax_computation", R"( // CHECK: #[[MAP:.*]] = affine_map<()[s0] -> (s0 * 16)> // CHECK-LABEL: tt.func @triton_fn( // CHECK-SAME: %[[P0:[^:]*]]: !tt.ptr {tt.divisibility = 16 : i32}, @@ -4527,7 +4527,7 @@ ENTRY e { } )"; TritonGemmConfig config(32, 32, 32, 1, 1, 1); - ASSERT_OK( + TF_ASSERT_OK( CreateTritonIrAndFileCheck(kHloText, config, EmitMatMul, "triton_dot", R"( CHECK: %[[INFINITY:.*]] = arith.constant dense<0x7F800000> : tensor<32x32xf32> CHECK: %[[C_MASK:.*]] = arith.constant dense<-65536> : tensor<32x32xi32> @@ -4858,7 +4858,7 @@ ENTRY e { } )"; TritonGemmConfig config(32, 32, 32, 1, 1, 1); - ASSERT_OK( + TF_ASSERT_OK( CreateTritonIrAndFileCheck(kHloText, config, EmitMatMul, "triton_dot", R"( CHECK: %[[INFINITY:.*]] = arith.constant dense<0x7F800000> : tensor<32x32xf32> CHECK: %[[C_MASK:.*]] = arith.constant dense<-65536> : tensor<32x32xi32>