diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc index 690887b5f83840..42d2deeca92700 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc @@ -676,6 +676,16 @@ CodegenDecision IsTritonSupportedDot( CodegenDecision IsTritonSupportedScaledDot( const HloScaledDotInstruction& dot, const se::GpuComputeCapability& gpu_version) { + if (gpu_version.IsCuda()) { + auto cc = gpu_version.cuda_compute_capability(); + if (!cc || !cc->IsAtLeastAmpere()) { + return CodegenDecision::Forbid( + "Scaled dot is not supported by Triton for pre-Ampere GPUs."); + } + } else if (!gpu_version.IsRocm()) { + return CodegenDecision::Forbid( + "Scaled dot is only supported on CUDA and ROCm."); + } CHECK_GE(dot.operand_count(), 4); PrimitiveType lhs_type = dot.operand(0)->shape().element_type(); PrimitiveType rhs_type = dot.operand(1)->shape().element_type(); diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc index 21592918b634f0..e30edaf4702c5f 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc @@ -2628,6 +2628,27 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Bool()), ScaledDotTestName); +TEST_F(HloHardwareIndependentTestBase, ScaledDotPreAmpereIsRejected) { + const std::string kHlo = R"( +HloModule ScaledDotPreAmpere + +ENTRY triton_computation { + lhs = f8e4m3fn[16, 32] parameter(0) + lhs_scale = f8e8m0fnu[16, 1] parameter(1) + rhs = f8e4m3fn[32, 16] parameter(2) + rhs_scale = f8e8m0fnu[1, 16] parameter(3) + ROOT dot = f32[16, 16] scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_contracting_dims={1}, + rhs_contracting_dims={0}, + backend_config={sizes:[16]} +} +)"; + TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo)); + const HloInstruction* dot = module->entry_computation()->root_instruction(); + EXPECT_FALSE(IsTritonSupportedInstruction( + *dot, se::GpuComputeCapability(se::CudaComputeCapability::Volta()))); +} + TEST_P(SupportTestWithTilingParam, NestedFusionsAreRejected) { // Nested fusions are not supported by xtile emitter. absl::string_view hlo_text = R"( diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD index ac4f5a34ea0178..e241fc166c2552 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD @@ -181,14 +181,17 @@ xla_test( "//xla/backends/gpu/codegen/triton:xtile_compiler", "//xla/backends/gpu/codegen/triton:xtile_test_base", "//xla/backends/gpu/tests:gpu_pjrt_codegen_test", + "//xla/backends/gpu/transforms:composite_rewriter", "//xla/codegen/xtile:block_level_parameters", "//xla/codegen/xtile:xtile_config_proto_cc", "//xla/hlo/ir:hlo", + "//xla/hlo/parser:hlo_parser", "//xla/hlo/testlib:filecheck", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/testlib:verified_hlo_module", "//xla/service:algorithm_util", "//xla/service/gpu:backend_configs_cc", + "//xla/service/gpu:gpu_compiler", "//xla/service/gpu:gpu_device_info_for_tests", "//xla/service/gpu:target_constants", "//xla/stream_executor:device_description", @@ -202,7 +205,9 @@ xla_test( "//xla/tsl/platform:errors", "//xla/tsl/platform:statusor", "//xla/tsl/platform:test", + "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/random:distributions", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", @@ -211,6 +216,7 @@ xla_test( "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest", "@eigen_archive//:eigen3", + "@llvm-project//llvm:Support", "@llvm-project//llvm:TargetParser", "@llvm-project//llvm:ir_headers", "@llvm-project//mlir:IR", diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc index 2918ff3956fe5f..e0e1be21c78b8b 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc @@ -15,7 +15,9 @@ limitations under the License. #include #include +#include #include +#include #include #include #include @@ -27,19 +29,24 @@ limitations under the License. #include #include +#include "absl/algorithm/container.h" +#include "absl/log/check.h" #include "absl/log/log.h" #include "absl/random/uniform_int_distribution.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/status_matchers.h" +#include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "absl/strings/str_replace.h" +#include "absl/strings/str_split.h" #include "absl/strings/string_view.h" #include "absl/strings/substitute.h" #include "absl/types/span.h" #include "Eigen/Core" #include "llvm/IR/LLVMContext.h" +#include "llvm/Support/raw_ostream.h" #include "llvm/TargetParser/Triple.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/MLIRContext.h" @@ -51,6 +58,7 @@ limitations under the License. #include "xla/backends/gpu/codegen/triton/xtile_compiler.h" #include "xla/backends/gpu/codegen/triton/xtile_test_base.h" #include "xla/backends/gpu/tests/gpu_pjrt_codegen_test.h" +#include "xla/backends/gpu/transforms/composite_rewriter.h" #include "xla/codegen/xtile/block_level_parameters.h" #include "xla/codegen/xtile/xtile_config.pb.h" #include "xla/error_spec.h" @@ -58,6 +66,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/parser/hlo_parser.h" #include "xla/hlo/testlib/filecheck.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/verified_hlo_module.h" @@ -66,6 +75,7 @@ limitations under the License. #include "xla/primitive_util.h" #include "xla/service/algorithm_util.h" #include "xla/service/gpu/backend_configs.pb.h" +#include "xla/service/gpu/gpu_compiler.h" #include "xla/service/gpu/gpu_device_info_for_tests.h" #include "xla/service/gpu/target_constants.h" #include "xla/shape.h" @@ -2592,6 +2602,7 @@ struct Fp4ScaledDotTypeCase { PrimitiveType lhs_type; PrimitiveType rhs_type; PrimitiveType scale_type; + int block_size = 32; }; using Fp4ScaledDotTestParam = std::tuple arguments, MakeScaledDotArguments(optimized_module.get())); + if (block_size == 16 && GetCudaComputeCapability().IsAtLeastBlackwell()) { + EXPECT_EXIT( + { + EXPECT_THAT(RunAndCompareNoHloPasses( + std::move(optimized_module), + LiteralUtil::MakePointers(arguments), + ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}) + .message(), + HasSubstr("CUDA_ERROR_INVALID_INSTRUCTION")); + std::_Exit(1); + }, + ::testing::ExitedWithCode(1), ""); + return; + } EXPECT_TRUE(RunAndCompareNoHloPasses( std::move(optimized_module), LiteralUtil::MakePointers(arguments), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3})); @@ -2693,27 +2739,29 @@ TEST_P(TritonFp4ScaledDotTest, Executes) { GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; } const Fp4ScaledDotTypeCase& type_case = std::get<0>(GetParam()); - RunFp4ScaledDotExecutionTest(type_case.lhs_type, type_case.rhs_type, - type_case.scale_type, std::get<1>(GetParam()), - std::get<2>(GetParam())); + RunFp4ScaledDotExecutionTest( + type_case.lhs_type, type_case.rhs_type, type_case.scale_type, + type_case.block_size, std::get<1>(GetParam()), std::get<2>(GetParam())); } INSTANTIATE_TEST_SUITE_P( TritonFp4ScaledDotTestSuite, TritonFp4ScaledDotTest, ::testing::Combine( ::testing::Values( - Fp4ScaledDotTypeCase{F4E2M1FN, F4E2M1FN, - F8E8M0FNU}, // MXFP4 x MXFP4 - Fp4ScaledDotTypeCase{F4E2M1FN, F8E4M3FN, - F8E8M0FNU}, // MXFP4 x MXFP8 (MMAv5 mxf8f6f4) - Fp4ScaledDotTypeCase{F8E4M3FN, F4E2M1FN, - F8E8M0FNU}, // MXFP8 x MXFP4 (MMAv5 mxf8f6f4) - Fp4ScaledDotTypeCase{F4E2M1FN, F8E5M2, - F8E8M0FNU}, // MXFP4 x MXE5M2 - Fp4ScaledDotTypeCase{F8E5M2, F4E2M1FN, - F8E8M0FNU}, // MXE5M2 x MXFP4 - Fp4ScaledDotTypeCase{F4E2M1FN, F4E2M1FN, - F8E4M3FN}), // NVFP4 x NVFP4 + Fp4ScaledDotTypeCase{F4E2M1FN, F4E2M1FN, F8E8M0FNU, + 32}, // MXFP4 x MXFP4 (block 32) + Fp4ScaledDotTypeCase{F4E2M1FN, F4E2M1FN, F8E8M0FNU, + 16}, // MXFP4 x MXFP4 (block 16) + Fp4ScaledDotTypeCase{F4E2M1FN, F8E4M3FN, F8E8M0FNU, + 32}, // MXFP4 x MXFP8 (MMAv5 mxf8f6f4) + Fp4ScaledDotTypeCase{F8E4M3FN, F4E2M1FN, F8E8M0FNU, + 32}, // MXFP8 x MXFP4 (MMAv5 mxf8f6f4) + Fp4ScaledDotTypeCase{F4E2M1FN, F8E5M2, F8E8M0FNU, + 32}, // MXFP4 x MXE5M2 + Fp4ScaledDotTypeCase{F8E5M2, F4E2M1FN, F8E8M0FNU, + 32}, // MXE5M2 x MXFP4 + Fp4ScaledDotTypeCase{F4E2M1FN, F4E2M1FN, F8E4M3FN, + 16}), // NVFP4 x NVFP4 (block 16) ::testing::Bool(), ::testing::Bool(), ::testing::Bool()), Fp4ScaledDotTestParamToString); @@ -3259,17 +3307,19 @@ class TritonScaledDotCoverageTest : public TritonScaledDotTestBase, std::vector GetCoverageTestCases() { std::vector input_types = {F8E4M3FN, F8E5M2, F4E2M1FN}; - std::vector scale_types = {F8E8M0FNU}; std::vector cases; + // 1. All input combinations with E8M0 scale and block size 32 (9 cases) for (auto lhs : input_types) { for (auto rhs : input_types) { - for (auto scale : scale_types) { - int block_size = 32; - cases.push_back({lhs, rhs, scale, block_size}); - } + cases.push_back({lhs, rhs, F8E8M0FNU, 32}); } } + + // 2. Block size 16 cases + cases.push_back({F4E2M1FN, F4E2M1FN, F8E8M0FNU, 16}); + cases.push_back({F4E2M1FN, F4E2M1FN, F8E4M3FN, 16}); + return cases; } @@ -3287,32 +3337,38 @@ std::string ScaledDotCoverageTestParamToString( } TEST_P(TritonScaledDotCoverageTest, Executes) { - auto cc = GpuComputeCapability().cuda_compute_capability(); - if (!cc || !cc->IsAtLeastHopper()) { - GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; - } - const auto& [param, lhs_k_minor, rhs_k_minor, tiling_enabled] = GetParam(); - // Blackboard doesn't support non-E8M0 scale in TMEM? - if (cc->IsAtLeastBlackwell() && param.scale_type != F8E8M0FNU) { - GTEST_SKIP() << "Blackwell TMEM scale expects E8M0 scale, skipping other " - "scales for now."; + auto cc = GpuComputeCapability().cuda_compute_capability(); + std::string device_name = "Unknown"; + if (cc) { + device_name = cc->IsAtLeastBlackwell() + ? "B200" + : (cc->IsAtLeastHopper() ? "H100" : "Pre-Hopper"); } - constexpr absl::string_view kHloTemplate = R"hlo( -HloModule m + std::string lhs_name = + primitive_util::LowercasePrimitiveTypeName(param.lhs_type); + std::string rhs_name = + primitive_util::LowercasePrimitiveTypeName(param.rhs_type); + std::string scale_name = + primitive_util::LowercasePrimitiveTypeName(param.scale_type); + std::string lhs_display = + lhs_k_minor ? lhs_name : absl::StrCat(lhs_name, ".T"); + std::string rhs_display = + rhs_k_minor ? absl::StrCat(rhs_name, ".T") : rhs_name; + std::string tiling_name = tiling_enabled ? "Symbolic" : "Experimental"; + + LOG(ERROR) << "Report Device: " << device_name; + LOG(ERROR) << "Report LHS: " << lhs_display; + LOG(ERROR) << "Report RHS: " << rhs_display; + LOG(ERROR) << "Report Scale: " << scale_name; + LOG(ERROR) << "Report BlockSize: " << param.block_size; + LOG(ERROR) << "Report Tiling: " << tiling_name; -ENTRY e { - lhs = $lhs_type[$lhs_shape] parameter(0) - rhs = $rhs_type[$rhs_shape] parameter(1) - lhs_scale = $scale_type[$lhs_scale_shape] parameter(2) - rhs_scale = $scale_type[$rhs_scale_shape] parameter(3) - ROOT dot = bf16[$output_shape] scaled-dot(lhs, rhs, lhs_scale, rhs_scale), - lhs_contracting_dims={$lhs_contracting_dim}, - rhs_contracting_dims={$rhs_contracting_dim} -} -)hlo"; + if (!cc || !cc->IsAtLeastHopper()) { + GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs."; + } constexpr int64_t m = 128; constexpr int64_t n = 128; @@ -3330,46 +3386,218 @@ ENTRY e { ? absl::StrCat(n, ",", scale_k) : absl::StrCat(scale_k, ",", n); - std::string hlo = absl::StrReplaceAll( - kHloTemplate, - {{"$lhs_type", - primitive_util::LowercasePrimitiveTypeName(param.lhs_type)}, - {"$rhs_type", - primitive_util::LowercasePrimitiveTypeName(param.rhs_type)}, - {"$scale_type", - primitive_util::LowercasePrimitiveTypeName(param.scale_type)}, - {"$lhs_shape", lhs_shape}, - {"$rhs_shape", rhs_shape}, - {"$lhs_scale_shape", lhs_scale_shape}, - {"$rhs_scale_shape", rhs_scale_shape}, - {"$output_shape", absl::StrCat(m, ",", n)}, - {"$lhs_contracting_dim", lhs_k_minor ? "1" : "0"}, - {"$rhs_contracting_dim", rhs_k_minor ? "1" : "0"}}); - - ASSERT_OK_AND_ASSIGN(auto optimized_module, GetOptimizedModule(hlo)); - LOG(INFO) << "Optimized HLO Module:\n" << optimized_module->ToString(); + // 1. Stage: After Composite + std::string after_composite = "N/A"; + constexpr absl::string_view kCompositeHloTemplate = R"hlo( +HloModule test_composite + +%xla.scaled_dot.1 { + %p0 = $lhs_type[$lhs_shape]{1,0} parameter(0) + %p1 = $rhs_type[$rhs_shape]{1,0} parameter(1) + %p2 = $scale_type[$lhs_scale_shape]{1,0} parameter(2) + %p3 = $scale_type[$rhs_scale_shape]{1,0} parameter(3) + ROOT %dummy = bf16[$output_shape]{1,0} constant({...}) +} + +ENTRY %main { + %lhs = $lhs_type[$lhs_shape]{1,0} parameter(0) + %rhs = $rhs_type[$rhs_shape]{1,0} parameter(1) + %lhs_scale = $scale_type[$lhs_scale_shape]{1,0} parameter(2) + %rhs_scale = $scale_type[$rhs_scale_shape]{1,0} parameter(3) + ROOT %call = bf16[$output_shape]{1,0} call(%lhs, %rhs, %lhs_scale, %rhs_scale), + to_apply=%xla.scaled_dot.1, + is_composite=true, + frontend_attributes={ + composite.attributes="{dimension_numbers=[[[$lhs_contracting_dim],[$rhs_contracting_dim]],[[],[]]]}", + composite.name="xla.scaled_dot", + composite.version="1" + } +} +)hlo"; - HloComputation* scaled_dot_computation = GetFirstComputationWithInstruction( - *optimized_module, HloOpcode::kScaledDot); + std::string composite_hlo = + absl::StrReplaceAll(kCompositeHloTemplate, + {{"$lhs_type", lhs_name}, + {"$rhs_type", rhs_name}, + {"$scale_type", scale_name}, + {"$lhs_shape", lhs_shape}, + {"$rhs_shape", rhs_shape}, + {"$lhs_scale_shape", lhs_scale_shape}, + {"$rhs_scale_shape", rhs_scale_shape}, + {"$output_shape", absl::StrCat(m, ",", n)}, + {"$lhs_contracting_dim", lhs_k_minor ? "1" : "0"}, + {"$rhs_contracting_dim", rhs_k_minor ? "1" : "0"}}); + + if (auto comp_module = ParseAndReturnUnverifiedModule(composite_hlo); + comp_module.ok()) { + CompositeRewriter rewriter; + if (auto rewrite_status = rewriter.Run(comp_module->get()); + rewrite_status.ok()) { + const auto* root = + (*comp_module)->entry_computation()->root_instruction(); + after_composite = HloOpcodeString(root->opcode()); + } + } + LOG(ERROR) << "Report AfterComposite: " << after_composite; - ASSERT_NE(scaled_dot_computation, nullptr) - << "Scaled dot instruction not found in optimized module."; + // 2. Stage: Optimized HLO + constexpr absl::string_view kHloTemplate = R"hlo( +HloModule m - EXPECT_THAT(CreateTritonIrAndFileCheckForDot(*scaled_dot_computation, - "CHECK: tt.dot_scaled"), - IsOk()); +ENTRY e { + lhs = $lhs_type[$lhs_shape] parameter(0) + rhs = $rhs_type[$rhs_shape] parameter(1) + lhs_scale = $scale_type[$lhs_scale_shape] parameter(2) + rhs_scale = $scale_type[$rhs_scale_shape] parameter(3) + ROOT dot = bf16[$output_shape] scaled-dot(lhs, rhs, lhs_scale, rhs_scale), + lhs_contracting_dims={$lhs_contracting_dim}, + rhs_contracting_dims={$rhs_contracting_dim} +} +)hlo"; + + std::string hlo = absl::StrReplaceAll( + kHloTemplate, {{"$lhs_type", lhs_name}, + {"$rhs_type", rhs_name}, + {"$scale_type", scale_name}, + {"$lhs_shape", lhs_shape}, + {"$rhs_shape", rhs_shape}, + {"$lhs_scale_shape", lhs_scale_shape}, + {"$rhs_scale_shape", rhs_scale_shape}, + {"$output_shape", absl::StrCat(m, ",", n)}, + {"$lhs_contracting_dim", lhs_k_minor ? "1" : "0"}, + {"$rhs_contracting_dim", rhs_k_minor ? "1" : "0"}}); + + if (param.scale_type == F8E8M0FNU && param.block_size == 16 && + GetCudaComputeCapability().IsAtLeastBlackwell() && !lhs_k_minor) { +#ifndef NDEBUG + EXPECT_DEATH( + { (void)GetOptimizedModule(hlo); }, + "MMAv5 with kind=mxf4nvf4 does not support transpose"); + return; +#endif + } - ASSERT_OK_AND_ASSIGN(std::vector arguments, - MakeScaledDotArguments(optimized_module.get())); + std::string optimized_hlo = "N/A"; + auto optimized_module_or = GetOptimizedModule(hlo); + HloComputation* scaled_dot_computation = nullptr; + if (optimized_module_or.ok()) { + scaled_dot_computation = GetFirstComputationWithInstruction( + **optimized_module_or, HloOpcode::kScaledDot); + if (scaled_dot_computation) { + optimized_hlo = "scaled-dot"; + } else if (GetFirstComputationWithInstruction(**optimized_module_or, + HloOpcode::kDot)) { + optimized_hlo = "dot"; + } + } + LOG(ERROR) << "Report OptimizedHLO: " << optimized_hlo; + + // 3. Stage: Triton Dot Lowering + std::string triton_dot = "N/A"; + if (scaled_dot_computation && scaled_dot_computation->FusionInstruction()) { + mlir::MLIRContext mlir_context; + BlockLevelParameters block_level_parameters; + if (auto gpu_config = scaled_dot_computation->FusionInstruction() + ->backend_config(); + gpu_config.ok() && gpu_config->has_fusion_backend_config() && + gpu_config->fusion_backend_config().has_block_level_fusion_config()) { + block_level_parameters = BlockLevelParameters::FromBlockLevelFusionConfig( + gpu_config->fusion_backend_config().block_level_fusion_config()); + } + auto triton_source_or = + CreateTritonModule("triton_fn", + *Cast( + scaled_dot_computation->FusionInstruction()), + TestGpuDeviceInfo::RTXA6000DeviceInfo(*cc), + block_level_parameters, mlir_context); + if (triton_source_or.ok()) { + std::string mlir_str; + llvm::raw_string_ostream os(mlir_str); + triton_source_or->module()->print(os); + if (absl::StrContains(mlir_str, "tt.dot_scaled")) { + triton_dot = "tt.dot_scaled"; + } else if (absl::StrContains(mlir_str, "tt.dot")) { + triton_dot = "tt.dot"; + } else if (absl::StrContains(mlir_str, "xtile.dot_scaled")) { + triton_dot = "xtile.dot_scaled"; + } + } + } + LOG(ERROR) << "Report TritonDot: " << triton_dot; + + // 4. Stage: PTX Generation + std::vector ptx_instructions; + GpuCompiler* gpu_compiler = dynamic_cast(compiler()); + bool compilation_succeeded = false; + if (gpu_compiler && optimized_module_or.ok()) { + gpu_compiler->SetAsmHook([&](absl::string_view ptx) { + std::string ptx_str(ptx); + std::string line; + size_t pos = 0; + while ((pos = ptx_str.find('\n')) != std::string::npos) { + line = ptx_str.substr(0, pos); + size_t first_non_space = line.find_first_not_of(" \t"); + if (first_non_space != std::string::npos) { + line = line.substr(first_non_space); + } + std::vector tokens = + absl::StrSplit(line, absl::ByAnyChar(" \t;"), absl::SkipEmpty()); + if (!tokens.empty()) { + std::string inst = tokens[0]; + if (absl::StartsWith(inst, "@") && tokens.size() > 1) { + inst = tokens[1]; + } + if (absl::StrContains(inst, "mma") && + (absl::StartsWith(inst, "wgmma.mma") || + absl::StartsWith(inst, "tcgen05.mma") || + absl::StartsWith(inst, "mma.sync"))) { + ptx_instructions.push_back(inst); + } + } + ptx_str.erase(0, pos + 1); + } + }); - // Assert that arguments (inputs and scales) are not all zeros - for (int i = 0; i < arguments.size(); ++i) { - EXPECT_FALSE(arguments[i].IsAll(0)) << "Argument " << i << " is all zeros"; + auto cloned_module = (*optimized_module_or)->Clone(); + auto executable_or = CompileToExecutable(std::move(cloned_module), + /*run_optimization_passes=*/true); + compilation_succeeded = executable_or.ok(); + gpu_compiler->RemoveAsmHook(); } - EXPECT_TRUE(RunAndCompareNoHloPasses(std::move(optimized_module), - LiteralUtil::MakePointers(arguments), - ErrorSpec{/*aabs=*/0.1, /*arel=*/0.1})); + std::vector unique_ptx; + for (const auto& inst : ptx_instructions) { + if (absl::c_find(unique_ptx, inst) == unique_ptx.end()) { + unique_ptx.push_back(inst); + } + } + std::string ptx_str = + unique_ptx.empty() ? "N/A" : absl::StrJoin(unique_ptx, ", "); + LOG(ERROR) << "Report PTX: " << ptx_str; + + // 5. Stage: Execution Status + std::string status = "FAILURE"; + if (compilation_succeeded) { + auto arguments_or = MakeScaledDotArguments((*optimized_module_or).get()); + if (arguments_or.ok()) { + auto run_result = + RunAndCompareNoHloPasses(std::move(*optimized_module_or), + LiteralUtil::MakePointers(*arguments_or), + ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}); + status = run_result ? "SUCCESS" : "FAILURE"; + if (!run_result) { + if (absl::StrContains(run_result.message(), "ILLEGAL_INSTRUCTION")) { + status = "ILLEGAL_INSTRUCTION"; + } + } + } + } + LOG(ERROR) << "Report Status: " << status; + if (GetCudaComputeCapability().IsAtLeastBlackwell() && + status == "ILLEGAL_INSTRUCTION") { + EXPECT_DEATH( + { CHECK_NE(status, "ILLEGAL_INSTRUCTION"); }, "ILLEGAL_INSTRUCTION"); + } } INSTANTIATE_TEST_SUITE_P( diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index f8107004477ca0..f825adfc9e98d7 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -3013,14 +3013,11 @@ cc_library( srcs = ["composite_rewriter.cc"], hdrs = ["composite_rewriter.h"], deps = [ - "//xla:literal", "//xla:util", "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", "//xla/service:decision", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", "@com_google_absl//absl/status", @@ -3042,7 +3039,6 @@ xla_cc_test( ":composite_rewriter", "//xla/hlo/ir:hlo", "//xla/hlo/parser:hlo_parser", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/log", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings", diff --git a/third_party/xla/xla/backends/gpu/transforms/composite_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/composite_rewriter.cc index 855fe918a7285e..8e85af83edeb3d 100644 --- a/third_party/xla/xla/backends/gpu/transforms/composite_rewriter.cc +++ b/third_party/xla/xla/backends/gpu/transforms/composite_rewriter.cc @@ -15,8 +15,6 @@ limitations under the License. #include "xla/backends/gpu/transforms/composite_rewriter.h" -#include - #include "absl/container/flat_hash_set.h" #include "absl/log/log.h" #include "absl/status/status.h" @@ -36,10 +34,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" -#include "xla/literal.h" #include "xla/service/decision.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -141,102 +136,6 @@ Decision HasSupportedDimensionNumbers( return Decision::Allow(); } -Decision IsValidBf16Scale(const HloInstruction* operand, - const HloInstruction* scale) { - if (scale->shape().dimensions().size() != - operand->shape().dimensions().size()) { - return Decision::Forbid("scale and operand rank mismatch for BF16"); - } - for (int64_t dim : scale->shape().dimensions()) { - if (dim != 1) { - return Decision::Forbid("scale dim != 1 for BF16"); - } - } - if (scale->opcode() != HloOpcode::kConstant) { - return Decision::Forbid("scale is not constant for BF16"); - } - if (!scale->literal().IsAllFloat(1.0)) { - return Decision::Forbid("scale is not 1.0 for BF16"); - } - return Decision::Allow(); -} - -Decision IsSupportedBf16(const HloInstruction* lhs, const HloInstruction* rhs, - const HloInstruction* lhs_scale, - const HloInstruction* rhs_scale) { - if (lhs_scale->shape().element_type() != BF16 || - rhs_scale->shape().element_type() != BF16) { - return Decision::Forbid("BF16 operands require BF16 scales"); - } - if (Decision d = IsValidBf16Scale(lhs, lhs_scale); !d) { - return d; - } - return IsValidBf16Scale(rhs, rhs_scale); -} - -Decision IsSupportedScaleAndOperand(const HloInstruction* operand, - const HloInstruction* scale, - int64_t contracting_dim) { - PrimitiveType op_type = operand->shape().element_type(); - PrimitiveType scale_type = scale->shape().element_type(); - - if (scale_type != F8E8M0FNU) { - return Decision::Forbid(absl::StrCat("Unsupported scale type: ", - PrimitiveType_Name(scale_type))); - } - - if (contracting_dim >= scale->shape().dimensions().size()) { - return Decision::Forbid("contracting_dim out of bounds for scale"); - } - int64_t operand_dim_size = operand->shape().dimensions(contracting_dim); - int64_t scale_dim_size = scale->shape().dimensions(contracting_dim); - - if (scale_dim_size == 0 || operand_dim_size % scale_dim_size != 0) { - return Decision::Forbid("operand_dim_size not divisible by scale_dim_size"); - } - int64_t scale_factor = operand_dim_size / scale_dim_size; - - if (op_type == F8E4M3FN || op_type == F8E5M2 || op_type == F4E2M1FN) { - if (scale_factor % 16 != 0) { - return Decision::Forbid( - absl::StrCat("scale_factor % 16 != 0: ", scale_factor)); - } - return Decision::Allow(); - } - - return Decision::Forbid( - absl::StrCat("Unsupported operand type: ", PrimitiveType_Name(op_type))); -} - -Decision IsSupportedScaledDot( - const HloCallInstruction* call, - const DotDimensionNumbers& dot_dimension_numbers) { - if (call->operand_count() < 4) { - return Decision::Forbid("Scaled dot call has fewer than 4 operands"); - } - const HloInstruction* lhs = call->operand(0); - const HloInstruction* rhs = call->operand(1); - const HloInstruction* lhs_scale = call->operand(2); - const HloInstruction* rhs_scale = call->operand(3); - - if (lhs->shape().element_type() == BF16 && - rhs->shape().element_type() == BF16) { - return IsSupportedBf16(lhs, rhs, lhs_scale, rhs_scale); - } - - int64_t lhs_contracting_dim = - dot_dimension_numbers.lhs_contracting_dimensions(0); - int64_t rhs_contracting_dim = - dot_dimension_numbers.rhs_contracting_dimensions(0); - - if (Decision d = - IsSupportedScaleAndOperand(lhs, lhs_scale, lhs_contracting_dim); - !d) { - return d; - } - return IsSupportedScaleAndOperand(rhs, rhs_scale, rhs_contracting_dim); -} - absl::Status RewriteScaledDot( HloComputation* computation, HloCallInstruction* call, const DotDimensionNumbers& dot_dimension_numbers) { @@ -251,7 +150,6 @@ absl::Status RewriteScaledDot( ABSL_RETURN_IF_ERROR(computation->RemoveInstruction(call)); return absl::OkStatus(); } - } // namespace absl::StatusOr CompositeRewriter::RewriteComputation( @@ -259,7 +157,7 @@ absl::StatusOr CompositeRewriter::RewriteComputation( bool changed = false; for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) { if (Decision d = IsScaledDotComposite(instruction); !d) { - VLOG(3) << d.Explain(); + VLOG(5) << d.Explain(); continue; } auto* call = Cast(instruction); @@ -277,11 +175,6 @@ absl::StatusOr CompositeRewriter::RewriteComputation( continue; } - if (Decision d = IsSupportedScaledDot(call, dot_dimension_numbers); !d) { - VLOG(3) << "Scaled dot composite operands not supported: " << d.Explain(); - continue; - } - ABSL_RETURN_IF_ERROR(RewriteScaledDot(computation, call, dot_dimension_numbers)); changed = true; } diff --git a/third_party/xla/xla/backends/gpu/transforms/composite_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/composite_rewriter_test.cc index efb285d3faa95d..bf717ab3e3ef83 100644 --- a/third_party/xla/xla/backends/gpu/transforms/composite_rewriter_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/composite_rewriter_test.cc @@ -27,7 +27,6 @@ limitations under the License. #include "absl/strings/substitute.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/parser/hlo_parser.h" -#include "xla/tsl/platform/statusor.h" namespace xla::gpu { namespace { @@ -162,8 +161,7 @@ TEST_P(CompositeRewriterParameterizedTest, Run) { LOG(INFO) << "HLO string: \n" << hlo_string; CompositeRewriter rewriter; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnUnverifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_string)); auto result = rewriter.Run(module.get()); @@ -172,117 +170,96 @@ TEST_P(CompositeRewriterParameterizedTest, Run) { EXPECT_THAT(module->entry_computation()->root_instruction()->opcode(), HloOpcode::kScaledDot); } else { - // If it didn't rewrite, it should either be OkAndHolds(false) - // or arguably just check that the opcode is still Call. - // The current implementation returns OkAndHolds(false) if no change. EXPECT_THAT(result, absl_testing::IsOkAndHolds(false)); EXPECT_THAT(module->entry_computation()->root_instruction()->opcode(), HloOpcode::kCall); } } +TEST(CompositeRewriterTest, MultipleContractingDimensionsRejects) { + absl::string_view hlo_string = R"( + HloModule test_module + + %xla.scaled_dot.1 { + %p0 = f8e4m3fn[128,256]{1,0} parameter(0) + %p1 = f8e4m3fn[256,128]{1,0} parameter(1) + %p2 = f8e8m0fnu[128,16]{1,0} parameter(2) + %p3 = f8e8m0fnu[16,128]{1,0} parameter(3) + ROOT %dummy = bf16[128,128]{1,0} constant({...}) + } + + ENTRY %main { + %lhs = f8e4m3fn[128,256]{1,0} parameter(0) + %rhs = f8e4m3fn[256,128]{1,0} parameter(1) + %lhs_scales = f8e8m0fnu[128,16]{1,0} parameter(2) + %rhs_scales = f8e8m0fnu[16,128]{1,0} parameter(3) + ROOT %call = bf16[128,128]{1,0} call(%lhs, %rhs, %lhs_scales, %rhs_scales), + to_apply=%xla.scaled_dot.1, + is_composite=true, + frontend_attributes={ + composite.attributes="{dimension_numbers=[[[1, 0],[1]],[[0],[0]]]}", + composite.name="xla.scaled_dot", + composite.version="1" + } + } + )"; + CompositeRewriter rewriter; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_string)); + EXPECT_THAT(rewriter.Run(module.get()), absl_testing::IsOkAndHolds(false)); +} + INSTANTIATE_TEST_SUITE_P( ScaledDotTests, CompositeRewriterParameterizedTest, - ::testing::Values( - TestCase{ - /*test_name=*/"FP8_Standard_Case", - /*lhs_type=*/"f8e4m3fn", - /*rhs_type=*/"f8e4m3fn", - /*lhs_scale_type=*/"f8e8m0fnu", - /*rhs_scale_type=*/"f8e8m0fnu", - /*lhs_scale_shape=*/"3,128,8", - /*rhs_scale_shape=*/"3,8,128", - /*lhs_scale_const_val=*/std::nullopt, - /*rhs_scale_const_val=*/std::nullopt, - /*expected_rewrite=*/true, - }, - TestCase{ - /*test_name=*/"BF16_Identity_Case", - /*lhs_type=*/"bf16", - /*rhs_type=*/"bf16", - /*lhs_scale_type=*/"bf16", - /*rhs_scale_type=*/"bf16", - /*lhs_scale_shape=*/"1,1,1", - /*rhs_scale_shape=*/"1,1,1", - /*lhs_scale_const_val=*/1.0f, - /*rhs_scale_const_val=*/1.0f, - /*expected_rewrite=*/true, - }, - TestCase{ - /*test_name=*/"BF16_Invalid_Scale_Value", - /*lhs_type=*/"bf16", - /*rhs_type=*/"bf16", - /*lhs_scale_type=*/"bf16", - /*rhs_scale_type=*/"bf16", - /*lhs_scale_shape=*/"1,1,1", - /*rhs_scale_shape=*/"1,1,1", - /*lhs_scale_const_val=*/1.0f, - /*rhs_scale_const_val=*/2.0f, - /*expected_rewrite=*/false, - }, - TestCase{ - /*test_name=*/"BF16_Invalid_Scale_Shape", - /*lhs_type=*/"bf16", - /*rhs_type=*/"bf16", - /*lhs_scale_type=*/"bf16", - /*rhs_scale_type=*/"bf16", - /*lhs_scale_shape=*/"3,128,1", - /*rhs_scale_shape=*/"1,1,1", - /*lhs_scale_const_val=*/std::nullopt, - /*rhs_scale_const_val=*/1.0f, - /*expected_rewrite=*/false, - }, - TestCase{ - /*test_name=*/"Mixed_Type_Fail_BF16_Scale_With_FP8_Op", - /*lhs_type=*/"f8e4m3fn", - /*rhs_type=*/"f8e4m3fn", - /*lhs_scale_type=*/"bf16", - /*rhs_scale_type=*/"f8e8m0fnu", - /*lhs_scale_shape=*/"3,128,8", - /*rhs_scale_shape=*/"3,8,128", - /*lhs_scale_const_val=*/std::nullopt, - /*rhs_scale_const_val=*/std::nullopt, - /*expected_rewrite=*/false, - }, - TestCase{ - /*test_name=*/"FP8_ScaleFactor_16", - /*lhs_type=*/"f8e4m3fn", - /*rhs_type=*/"f8e4m3fn", - /*lhs_scale_type=*/"f8e8m0fnu", - /*rhs_scale_type=*/"f8e8m0fnu", - /*lhs_scale_shape=*/"3,128,16", // 256 / 16 = 16 (divisible by 16) - /*rhs_scale_shape=*/"3,8,128", - /*lhs_scale_const_val=*/std::nullopt, - /*rhs_scale_const_val=*/std::nullopt, - /*expected_rewrite=*/true, - }, - TestCase{ - /*test_name=*/"FP8_ScaleFactor_Invalid", - /*lhs_type=*/"f8e4m3fn", - /*rhs_type=*/"f8e4m3fn", - /*lhs_scale_type=*/"f8e8m0fnu", - /*rhs_scale_type=*/"f8e8m0fnu", - /*lhs_scale_shape=*/"3,128,32", // K=256, scale_k=32 -> block_size - // = 8 (not divisible by 16) - /*rhs_scale_shape=*/"3,8,128", - /*lhs_scale_const_val=*/std::nullopt, - /*rhs_scale_const_val=*/std::nullopt, - /*expected_rewrite=*/false, - }, - TestCase{ - /*test_name=*/"FP8_ScaleFactor_64", - /*lhs_type=*/"f8e4m3fn", - /*rhs_type=*/"f8e4m3fn", - /*lhs_scale_type=*/"f8e8m0fnu", - /*rhs_scale_type=*/"f8e8m0fnu", - /*lhs_scale_shape=*/"3,128,4", // 256 / 4 = 64 (divisible by 32) - /*rhs_scale_shape=*/"3,8,128", - /*lhs_scale_const_val=*/std::nullopt, - /*rhs_scale_const_val=*/std::nullopt, - /*expected_rewrite=*/true, - }), + ::testing::Values(TestCase{ + /*test_name=*/"FP8_Standard_Case", + /*lhs_type=*/"f8e4m3fn", + /*rhs_type=*/"f8e4m3fn", + /*lhs_scale_type=*/"f8e8m0fnu", + /*rhs_scale_type=*/"f8e8m0fnu", + /*lhs_scale_shape=*/"3,128,8", + /*rhs_scale_shape=*/"3,8,128", + /*lhs_scale_const_val=*/std::nullopt, + /*rhs_scale_const_val=*/std::nullopt, + /*expected_rewrite=*/true, + }, + TestCase{ + /*test_name=*/"BF16_Identity_Case", + /*lhs_type=*/"bf16", + /*rhs_type=*/"bf16", + /*lhs_scale_type=*/"bf16", + /*rhs_scale_type=*/"bf16", + /*lhs_scale_shape=*/"1,1,1", + /*rhs_scale_shape=*/"1,1,1", + /*lhs_scale_const_val=*/1.0f, + /*rhs_scale_const_val=*/1.0f, + /*expected_rewrite=*/true, + }, + TestCase{ + /*test_name=*/"FP8_ScaleFactor_16", + /*lhs_type=*/"f8e4m3fn", + /*rhs_type=*/"f8e4m3fn", + /*lhs_scale_type=*/"f8e8m0fnu", + /*rhs_scale_type=*/"f8e8m0fnu", + /*lhs_scale_shape=*/"3,128,16", + /*rhs_scale_shape=*/"3,16,128", + /*lhs_scale_const_val=*/std::nullopt, + /*rhs_scale_const_val=*/std::nullopt, + /*expected_rewrite=*/true, + }, + TestCase{ + /*test_name=*/"FP4_E4M3_Scale_Block16", + /*lhs_type=*/"f4e2m1fn", + /*rhs_type=*/"f4e2m1fn", + /*lhs_scale_type=*/"f8e4m3fn", + /*rhs_scale_type=*/"f8e4m3fn", + /*lhs_scale_shape=*/"3,128,16", + /*rhs_scale_shape=*/"3,16,128", + /*lhs_scale_const_val=*/std::nullopt, + /*rhs_scale_const_val=*/std::nullopt, + /*expected_rewrite=*/true, + }), [](const ::testing::TestParamInfo& info) { return info.param.test_name; }); diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc index 424401adebb923..94e4cb7b6b7bb2 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc @@ -1754,6 +1754,7 @@ bool GpuCompiler::IsScaledDotSupportedByBackend( const se::GpuComputeCapability& gpu_version = gpu_target_config.device_description.gpu_compute_capability(); return debug_options.xla_gpu_experimental_scaled_dot_with_triton() && + IsTritonGemmEnabled(debug_options, gpu_version) && IsTritonSupportedInstruction(*instr, gpu_version).IsAllowed(); }