Skip to content
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

Jump to bottom

Some minor fixes #1297

Merged
merged 7 commits into from
May 24, 2024
Merged

Some minor fixes #1297

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
4d2a04c
support async collective op execution
eedalong Mar 14, 2024
f1039be
rebase master branch
eedalong Mar 22, 2024
e6b6339
merge master
eedalong May 20, 2024
5a6321c
merge master
eedalong May 20, 2024
e2a91b4
Merge remote-tracking branch 'origin/main' into support_async_collective
eedalong May 22, 2024
6613d88
Merge remote-tracking branch 'origin/main' into support_async_collective
eedalong May 23, 2024
6054c53
minor fixes
eedalong May 23, 2024
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
4 changes: 2 additions & 2 deletions .github/workflows/pytorch113_gpu.yml
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ jobs:
cuda_version: cu116
runner_tag: gpu-a10
remote_runtime_docker: bladedisc:latest-runtime-torch1.13.1-cu116
develop_base_image: nvidia/cuda:11.6.0-cudnn8-devel-ubuntu20.04
runtime_base_image: nvidia/cuda:11.6.0-cudnn8-devel-ubuntu20.04
develop_base_image: nvidia/cuda:11.6.1-cudnn8-devel-ubuntu20.04
runtime_base_image: nvidia/cuda:11.6.1-cudnn8-devel-ubuntu20.04
extra_build_args: --build-arg PYTHON_VERSION=PYTHON3.8 --build-arg ENABLE_FIND_FASTEST_APT_SOURCE=OFF
extra_envs: -e TORCH_BLADE_BUILD_TENSORRT_STATIC=OFF
-e TORCH_BLADE_CI_BUILD_TORCH_VERSION=1.13.1+cu116
Expand Down
4 changes: 2 additions & 2 deletions tao_compiler/mlir/disc/transforms/disc_op_schedule.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -354,7 +354,7 @@ class ScheduleGraph {
explicit ScheduleGraph(std::vector<Operation*>& post_order_instructions,
LatencyEstimator* latency_estimator,
AsyncTracker* async_tracker) {
InitilizeGrpahTopology(post_order_instructions, latency_estimator,
InitilizeGraphTopology(post_order_instructions, latency_estimator,
async_tracker);
InitializeGraphAnalysis(latency_estimator, async_tracker);
}
Expand Down Expand Up @@ -497,7 +497,7 @@ class ScheduleGraph {
}
}

void InitilizeGrpahTopology(std::vector<Operation*>& post_order_instructions,
void InitilizeGraphTopology(std::vector<Operation*>& post_order_instructions,
LatencyEstimator* latency_estimator,
AsyncTracker* async_tracker) {
original_order_ = post_order_instructions;
Expand Down
6 changes: 4 additions & 2 deletions tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5712,9 +5712,11 @@ struct DiscLhloLegalizeRootsToParallelLoops
// TODO: We should put even single nodes into a fusion by fusion pass
// Revisit this and walk lmhlo::FusionOp only after the revision done.
func.walk([&](lmhlo::LmhloOp op) {
// Skip the embedded ops in lmhlo.fusion or lmhlo.reduce/scatter
// Skip the embedded ops in lmhlo.fusion or lmhlo.reduce/scatter or
// lmhlo_disc.args_mutation
lmhlo::LmhloOp parent = op->getParentOfType<lmhlo::LmhloOp>();
if (parent && !isa<lmhlo::FusionOp>(op)) {
if (isa<lmhlo_disc::ArgsMutationOp>(op) ||
parent && !isa<lmhlo::FusionOp>(op)) {
return;
}
if (isFusionType<FusionType::kStitch>(op) &&
Expand Down
2 changes: 1 addition & 1 deletion tao_compiler/mlir/disc/transforms/mhlo_placer.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -418,7 +418,7 @@ void OpsPlacer::placeI32Ops() {
if (isa<mhlo_disc::CustomCallV2Op>(op)) return;

if (isa<mhlo::TupleOp, mhlo::GetTupleElementOp, mhlo::WhileOp, mhlo::IfOp,
mhlo::ReturnOp>(op)) {
mhlo::ReturnOp, mhlo_disc::ArgsMutationOp>(op)) {
return;
}
// Skip the Op that is already placed on CPU
Expand Down
Loading