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

Add depth=-1 mode to air-par-to-herd/segment/launch #866

Merged
merged 3 commits into from
Jan 20, 2025
Merged

Add depth=-1 mode to air-par-to-herd/segment/launch #866

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
bdc2b7f
Add depth=-1 mode to air-par-to-herd/segment/launch representing conv…
erwei-xilinx Jan 20, 2025
f45060b
Merge branch 'main' of github.com:erwei-xilinx/mlir-air-erwei into de…
erwei-xilinx Jan 20, 2025
d9dadb8
Enable depth=-1 for forall
erwei-xilinx Jan 20, 2025
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
18 changes: 12 additions & 6 deletions mlir/include/air/Conversion/Passes.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,10 @@ def ParallelToHerd : Pass<"air-par-to-herd", "ModuleOp"> {
}];
let options = [
Option<"clAssignDepth", "depth", "int",
/*default=*/"-1",
"Given a nest of parallel for loops, which depth to map to air.herd">,
/*default=*/"-2",
"Given a nest of parallel for loops, which depth to map to air.herd. "
"-1 means converting the innermost parallel loop; any other negative "
"value means converting all parallel loops">,
Option<"clFirstDim", "first-dim", "int",
/*default=*/"0",
"Which herd dimension to map to first. Can be zero or one. If set to "
Expand All @@ -49,8 +51,10 @@ def ParallelToLaunch : Pass<"air-par-to-launch", "ModuleOp"> {
}];
let options = [
Option<"clAssignDepth", "depth", "int",
/*default=*/"-1",
"Given a nest of parallel for loops, which depth to map to air.launch">,
/*default=*/"-2",
"Given a nest of parallel for loops, which depth to map to air.launch"
"-1 means converting the innermost parallel loop; any other negative "
"value means converting all parallel loops">,
Option<"clHasSegment", "has-air-segment", "bool", /*default=*/"false",
"Whether to create an air.segment op in generated air.launch "
"regions">,
Expand All @@ -68,8 +72,10 @@ def ParallelToSegment : Pass<"air-par-to-segment", "ModuleOp"> {
}];
let options = [
Option<"clAssignDepth", "depth", "int",
/*default=*/"-1",
"Given a nest of parallel for loops, which depth to map to air.segment">,
/*default=*/"-2",
"Given a nest of parallel for loops, which depth to map to air.segment"
"-1 means converting the innermost parallel loop; any other negative "
"value means converting all parallel loops">,
];
}

Expand Down
48 changes: 45 additions & 3 deletions mlir/lib/Conversion/ConvertToAIRPass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1165,11 +1165,25 @@ struct ParallelToHerdPass
if (llvm::any_of(hierOps,
[op](Operation *h) { return op->isProperAncestor(h); }))
return;
// Depth = -1 means converting the innermost parallel ops
if (clAssignDepth == -1) {
SmallVector<Operation *> parOpsInOp;
op->walk([&parOpsInOp](Operation *o) {
if (isa<scf::ForallOp, scf::ParallelOp, affine::AffineParallelOp>(o))
parOpsInOp.push_back(o);
});
if (parOpsInOp.size() > 1)
return;
filteredOps.insert(op);
return;
}
// Assigning depth to other negative values means converting all
// parallel ops
if (clAssignDepth < 0) {
filteredOps.insert(op);
return;
}
// the number of nested scf.parallel above this one
// the number of nested parallel above this one
int parallel_depth = 0;
Operation *par = op;
while ((par = par->getParentOp()))
Expand Down Expand Up @@ -1253,11 +1267,25 @@ struct ParallelToLaunchPass
return op->isProperAncestor(l);
}))
return;
// Depth = -1 means converting the innermost parallel ops
if (clAssignDepth == -1) {
SmallVector<Operation *> parOpsInOp;
op->walk([&parOpsInOp](Operation *o) {
if (isa<scf::ParallelOp>(o))
parOpsInOp.push_back(o);
});
if (parOpsInOp.size() > 1)
return;
filteredOps.insert(op);
return;
}
// Assigning depth to other negative values means converting all
// parallel ops
if (clAssignDepth < 0) {
filteredOps.insert(op);
return;
}
// the number of nested scf.parallel above this one
// the number of nested parallel above this one
int parallel_depth = 0;
Operation *par = op;
while ((par = par->getParentOp()))
Expand Down Expand Up @@ -1342,11 +1370,25 @@ struct ParallelToSegmentPass
return op->isProperAncestor(s);
}))
return;
// Depth = -1 means converting the innermost parallel ops
if (clAssignDepth == -1) {
SmallVector<Operation *> parOpsInOp;
op->walk([&parOpsInOp](Operation *o) {
if (isa<scf::ParallelOp>(o))
parOpsInOp.push_back(o);
});
if (parOpsInOp.size() > 1)
return;
filteredOps.insert(op);
return;
}
// Assigning depth to other negative values means converting all
// parallel ops
if (clAssignDepth < 0) {
filteredOps.insert(op);
return;
}
// the number of nested scf.parallel above this one
// the number of nested parallel above this one
int parallel_depth = 0;
Operation *par = op;
while ((par = par->getParentOp()))
Expand Down
24 changes: 24 additions & 0 deletions mlir/test/Conversion/ConvertToAIR/scf_forall_to_herd.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@
//===----------------------------------------------------------------------===//

// RUN: air-opt -split-input-file -verify-diagnostics -air-par-to-herd %s | FileCheck %s
// RUN: air-opt -split-input-file -verify-diagnostics -air-par-to-herd="depth=-1" %s | FileCheck %s --check-prefix=DEPTHM1
// RUN: air-opt -split-input-file -verify-diagnostics -air-par-to-herd="depth=0" %s | FileCheck %s --check-prefix=DEPTH0
// RUN: air-opt -split-input-file -verify-diagnostics -air-par-to-herd="depth=1" %s | FileCheck %s --check-prefix=DEPTH1

// CHECK-LABEL: func.func @scf0() {
// CHECK: air.herd @herd_0 tile (%{{.*}}, %{{.*}}) in (%{{.*}}=%c2{{.*}}, %{{.*}}=%c2{{.*}})
Expand Down Expand Up @@ -98,6 +101,27 @@ func.func @scf4() {
// CHECK: }
// CHECK: }
// CHECK: }
// DEPTHM1-LABEL: func.func @scf5() {
// DEPTHM1: scf.forall {{.*}} {
// DEPTHM1: scf.forall {{.*}} {
// DEPTHM1: air.herd @herd_{{.*}} {
// DEPTHM1: }
// DEPTHM1: }
// DEPTHM1: }
// DEPTH0-LABEL: func.func @scf5() {
// DEPTH0: air.herd @herd_{{.*}} {
// DEPTH0: scf.forall {{.*}} {
// DEPTH0: scf.forall {{.*}} {
// DEPTH0: }
// DEPTH0: }
// DEPTH0: }
// DEPTH1-LABEL: func.func @scf5() {
// DEPTH1: scf.forall {{.*}} {
// DEPTH1: air.herd @herd_{{.*}} {
// DEPTH1: scf.forall {{.*}} {
// DEPTH1: }
// DEPTH1: }
// DEPTH1: }
func.func @scf5() {
%src = memref.alloc() : memref<4x4x4xi32, 2 : i32>
%dst = memref.alloc() : memref<4x4x4xi32, 2 : i32>
Expand Down
76 changes: 76 additions & 0 deletions mlir/test/Conversion/ConvertToAIR/scf_parallel_to_herd.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
//===----------------------------------------------------------------------===//

// RUN: air-opt -split-input-file -verify-diagnostics -air-par-to-herd %s | FileCheck %s
// RUN: air-opt -split-input-file -verify-diagnostics -air-par-to-herd="depth=-1" %s | FileCheck %s --check-prefix=DEPTHM1
// RUN: air-opt -split-input-file -verify-diagnostics -air-par-to-herd="depth=0" %s | FileCheck %s --check-prefix=DEPTH0

// CHECK-LABEL: func.func @scf0() {
// CHECK: %[[C2:.*]] = arith.constant 2 : index
Expand Down Expand Up @@ -201,6 +203,34 @@ module {
// CHECK: return
// CHECK: }
// CHECK: }
// DEPTHM1-LABEL: @shared_herd_name
// DEPTHM1: scf.parallel {{.*}} {
// DEPTHM1: air.herd @herd_0
// DEPTHM1: }
// DEPTHM1: air.herd @herd_0
// DEPTHM1: }
// DEPTHM1: air.herd @herd_0
// DEPTHM1: }
// DEPTHM1: scf.reduce
// DEPTHM1: }
// DEPTHM1: return
// DEPTHM1: }
// DEPTHM1: }
// DEPTH0-LABEL: @shared_herd_name
// DEPTH0: air.herd @herd_0
// DEPTH0: scf.parallel {{.*}}
// DEPTH0: scf.reduce
// DEPTH0: }
// DEPTH0: scf.parallel {{.*}}
// DEPTH0: scf.reduce
// DEPTH0: }
// DEPTH0: scf.parallel {{.*}}
// DEPTH0: scf.reduce
// DEPTH0: }
// DEPTH0: }
// DEPTH0: return
// DEPTH0: }
// DEPTH0: }
module {
func.func @shared_herd_name(%arg0: memref<512x1024xbf16>, %arg1: memref<1024x512xbf16>, %arg2: memref<512x512xbf16>) {
%c32 = arith.constant 32 : index
Expand Down Expand Up @@ -248,6 +278,29 @@ module {
// CHECK: return
// CHECK: }
// CHECK: }
// DEPTHM1-LABEL: @unique_herd_name
// DEPTHM1: scf.parallel {{.*}} {
// DEPTHM1: air.herd @herd_0
// DEPTHM1: }
// DEPTHM1: air.herd @herd_1
// DEPTHM1: }
// DEPTHM1: scf.reduce
// DEPTHM1: }
// DEPTHM1: return
// DEPTHM1: }
// DEPTHM1: }
// DEPTH0-LABEL: @unique_herd_name
// DEPTH0: air.herd @herd_0
// DEPTH0: scf.parallel {{.*}} {
// DEPTH0: scf.reduce
// DEPTH0: }
// DEPTH0: scf.parallel {{.*}} {
// DEPTH0: scf.reduce
// DEPTH0: }
// DEPTH0: }
// DEPTH0: return
// DEPTH0: }
// DEPTH0: }
module {
func.func @unique_herd_name(%arg0: memref<512x1024xbf16>, %arg1: memref<1024x512xbf16>, %arg2: memref<512x512xbf16>) {
%c32 = arith.constant 32 : index
Expand Down Expand Up @@ -303,6 +356,29 @@ module {
// CHECK: return
// CHECK: }
// CHECK: }
// DEPTHM1-LABEL: @l2_to_l1_dma_infer_herd
// DEPTHM1: scf.parallel {{.*}} {
// DEPTHM1: air.herd @herd_0
// DEPTHM1: }
// DEPTHM1: air.herd @herd_0
// DEPTHM1: }
// DEPTHM1: scf.reduce
// DEPTHM1: }
// DEPTHM1: return
// DEPTHM1: }
// DEPTHM1: }
// DEPTH0-LABEL: @l2_to_l1_dma_infer_herd
// DEPTH0: air.herd @herd_0
// DEPTH0: scf.parallel {{.*}} {
// DEPTH0: scf.reduce
// DEPTH0: }
// DEPTH0: scf.parallel {{.*}} {
// DEPTH0: scf.reduce
// DEPTH0: }
// DEPTH0: }
// DEPTH0: return
// DEPTH0: }
// DEPTH0: }
module {
func.func @l2_to_l1_dma_infer_herd() {
%c32 = arith.constant 32 : index
Expand Down
Loading