Skip to content

[XLA:MGPU][Experimental] HLO -> Pallas. #97357

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 24, 2025
Merged

[XLA:MGPU][Experimental] HLO -> Pallas. #97357

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions third_party/xla/xla/backends/gpu/codegen/triton/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -408,6 +408,7 @@ cc_library(
"//xla/service/gpu:launch_dimensions",
"//xla/service/gpu:matmul_utils",
"//xla/service/gpu:triton_fusion_analysis",
"//xla/service/gpu/model:symbolic_tile_analysis",
"//xla/service/gpu/model:tiled_hlo_instruction_or_computation",
"//xla/stream_executor:device_description",
"//xla/stream_executor:launch_dim",
Expand Down
28 changes: 11 additions & 17 deletions third_party/xla/xla/backends/gpu/codegen/triton/fusion_emitter.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1447,21 +1447,8 @@ absl::StatusOr<ScalarOrTensor> EmitScope(
}
} // namespace

namespace {

using ::xla::gpu::ir_emitter_triton_internal::DumpTritonIR;
namespace ir_emitter_triton_internal {

// Given a tiling specification for a fusion and an annotated fusion, derives a
// tiling for the annotated fusion.
//
// Note that the tiling extracted here is voluntarily not checked against the
// specification, which means that it could be invalid. This should only be the
// case, though, if this logic gets stale, or if the fusion does not contain
// the required annotations. Checking constraints is not cheap, so we left it up
// to the caller to decide when to check the constraints.
//
// TODO(b/421837868): this belongs near/in `BlockLevelParameters`, but we start
// with this here in order to allow an incremental replacement.
absl::StatusOr<Tiling> TilingFromAnnotatedFusion(
const HloFusionInstruction* fusion,
const SymbolicTileAnalysis& symbolic_tile_analysis,
Expand Down Expand Up @@ -1514,6 +1501,12 @@ absl::StatusOr<Tiling> TilingFromAnnotatedFusion(
return Tiling(std::move(tile_mapping));
}

} // namespace ir_emitter_triton_internal

namespace {

using ::xla::gpu::ir_emitter_triton_internal::DumpTritonIR;

// Generate Triton IR inside 'fn', using the given block_level_parameters.
// TODO(b/421837868): `BlockLevelParameters` should hold all the necessary
// tiling information.
Expand Down Expand Up @@ -1542,9 +1535,10 @@ absl::StatusOr<SmallVector<Value>> EmitGeneric(

// TODO(b/421837868): unify the logic to extract tiling parameters with
// `BlockLevelParameters`.
TF_ASSIGN_OR_RETURN(Tiling tiling,
TilingFromAnnotatedFusion(fusion, symbolic_tile_analysis,
block_level_parameters));
TF_ASSIGN_OR_RETURN(
Tiling tiling,
ir_emitter_triton_internal::TilingFromAnnotatedFusion(
fusion, symbolic_tile_analysis, block_level_parameters));

// TODO(b/372454662): Decide which root to use. Currently, we only support
// "simple" multi-output fusions that have just one root without users. This
Expand Down
19 changes: 18 additions & 1 deletion third_party/xla/xla/backends/gpu/codegen/triton/fusion_emitter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ limitations under the License.
#include "xla/codegen/emitter_loc_op_builder.h"
#include "xla/hlo/ir/hlo_instructions.h"
#include "xla/hlo/ir/hlo_module.h"
#include "xla/service/gpu/model/symbolic_tile_analysis.h"
#include "xla/service/gpu/model/tiled_hlo_computation.h"
#include "xla/service/gpu/model/tiled_hlo_instruction.h"
#include "xla/service/hlo_module_config.h"
Expand Down Expand Up @@ -103,7 +104,7 @@ std::string GetLibdevicePath(const HloModuleConfig& hlo_config,

// TODO(b/406472229): Move the contents of this namespace to a helpers file
// to avoid polluting `fusion_emitter.h`.
// Exposed for testing purposes only. Do not use.
// Exposed for testing and experimental purposes only. Do not use.
namespace ir_emitter_triton_internal {

// Computes the transformation from a 1-d program_id to a tile multi-index.
Expand All @@ -127,6 +128,22 @@ inline std::string DumpTritonIR(mlir::ModuleOp triton_module,
return triton_ir;
}

// Given a tiling specification for a fusion and an annotated fusion, derives a
// tiling for the annotated fusion.
//
// Note that the tiling extracted here is voluntarily not checked against the
// specification, which means that it could be invalid. This should only be the
// case, though, if this logic gets stale, or if the fusion does not contain
// the required annotations. Checking constraints is not cheap, so we left it up
// to the caller to decide when to check the constraints.
//
// TODO(b/421837868): this belongs near/in `BlockLevelParameters`, but we start
// with this here in order to allow an incremental replacement.
absl::StatusOr<Tiling> TilingFromAnnotatedFusion(
const HloFusionInstruction* fusion,
const SymbolicTileAnalysis& symbolic_tile_analysis,
const BlockLevelParameters& block_level_parameters);

} // namespace ir_emitter_triton_internal
} // namespace gpu
} // namespace xla
Expand Down
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy