Skip to content

[aster] Revamp control-flow handling #489

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

Draft
fabianmcg wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[aster] Revamp control-flow handling #489

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
87 changes: 18 additions & 69 deletions include/aster/Dialect/AMDGCN/IR/AMDGCNTypes.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,14 +50,9 @@ def AGPRType : AMDGCN_RegisterDef<"AGPR", "agpr", [MemRefElementTypeInterface]>
let genVerifyDecl = 1;
let extraClassDeclaration = [{
Register getReg() const { return getRange().begin(); }

/// Returns true if the register is relocatable.
bool isRelocatable() const { return getReg().isRelocatable(); }

//===------------------------------------------------------------------===//
// RegisterTypeInterface
//===------------------------------------------------------------------===//
bool isRegisterRange() const { return getRange().size() > 1; }
RegisterRange getAsRange() const {
return getRange();
}
Expand All @@ -67,9 +62,6 @@ def AGPRType : AMDGCN_RegisterDef<"AGPR", "agpr", [MemRefElementTypeInterface]>
RegisterTypeInterface cloneRegisterType(RegisterRange range) const {
return AGPRType::get(getContext(), range);
}
RegisterTypeInterface cloneRegisterType(Register reg) const {
return AGPRType::get(getContext(), RegisterRange(reg, 1));
}

//===------------------------------------------------------------------===//
// ResourceTypeInterface
Expand Down Expand Up @@ -97,14 +89,9 @@ def SGPRType : AMDGCN_RegisterDef<"SGPR", "sgpr", [MemRefElementTypeInterface]>
let genVerifyDecl = 1;
let extraClassDeclaration = [{
Register getReg() const { return getRange().begin(); }

/// Returns true if the register is relocatable.
bool isRelocatable() const { return getReg().isRelocatable(); }

//===------------------------------------------------------------------===//
// RegisterTypeInterface
//===------------------------------------------------------------------===//
bool isRegisterRange() const { return getRange().size() > 1; }
RegisterRange getAsRange() const {
return getRange();
}
Expand All @@ -114,9 +101,6 @@ def SGPRType : AMDGCN_RegisterDef<"SGPR", "sgpr", [MemRefElementTypeInterface]>
RegisterTypeInterface cloneRegisterType(RegisterRange range) const {
return SGPRType::get(getContext(), range);
}
RegisterTypeInterface cloneRegisterType(Register reg) const {
return SGPRType::get(getContext(), RegisterRange(reg, 1));
}

//===------------------------------------------------------------------===//
// ResourceTypeInterface
Expand Down Expand Up @@ -145,13 +129,9 @@ def VGPRType : AMDGCN_RegisterDef<"VGPR", "vgpr", [MemRefElementTypeInterface]>
let extraClassDeclaration = [{
Register getReg() const { return getRange().begin(); }

/// Returns true if the register is relocatable.
bool isRelocatable() const { return getRange().begin().isRelocatable(); }

//===------------------------------------------------------------------===//
// RegisterTypeInterface
//===------------------------------------------------------------------===//
bool isRegisterRange() const { return getRange().size() > 1; }
RegisterRange getAsRange() const {
return getRange();
}
Expand All @@ -161,9 +141,6 @@ def VGPRType : AMDGCN_RegisterDef<"VGPR", "vgpr", [MemRefElementTypeInterface]>
RegisterTypeInterface cloneRegisterType(RegisterRange range) const {
return VGPRType::get(getContext(), range);
}
RegisterTypeInterface cloneRegisterType(Register reg) const {
return VGPRType::get(getContext(), RegisterRange(reg, 1));
}

//===------------------------------------------------------------------===//
// ResourceTypeInterface
Expand All @@ -176,71 +153,43 @@ def VGPRType : AMDGCN_RegisterDef<"VGPR", "vgpr", [MemRefElementTypeInterface]>
// SREG like types
//===----------------------------------------------------------------------===//

/// Special registers to model state.
def SREGType : AMDGCN_RegisterDef<"SREG", "sreg", [MemRefElementTypeInterface]> {
let summary = "SREG type";
let parameters = (ins
DefaultValuedParameter<"Register", "Register()">:$reg,
"SregKind":$kind);
let assemblyFormat = "`<`$kind (`,` $reg^)?`>`";
let genVerifyDecl = 1;
let extraClassDeclaration = [{
/// Returns true if the register is relocatable.
bool isRelocatable() const { return getReg().isRelocatable(); }

//===------------------------------------------------------------------===//
// RegisterTypeInterface
//===------------------------------------------------------------------===//
bool isRegisterRange() const { return false; }
RegisterRange getAsRange() const {
return RegisterRange(getReg(), 1);
}
RegisterKind getRegisterKind() const {
return RegisterKind::SREG;
}
RegisterTypeInterface cloneRegisterType(RegisterRange range) const {
assert(range.size() == 1 && "SREG type can only clone single register");
return SREGType::get(getContext(), range.begin(), getKind());
}
RegisterTypeInterface cloneRegisterType(Register reg) const {
return SREGType::get(getContext(), reg, getKind());
}

//===------------------------------------------------------------------===//
// ResourceTypeInterface
//===------------------------------------------------------------------===//
Resource *getResource() const;
}];
}

/// Special registers to model state.
class SREGBase<string name, string mnemonic, string kind>
: AMDGCN_RegisterDef<name, mnemonic, [MemRefElementTypeInterface]> {
let summary = kind # " special register type";
let assemblyFormat = "";
let parameters = (ins
DefaultValuedParameter<"Register", "Register()">:$reg
);
let assemblyFormat = "(`<` $reg^ `>`)?";
let builders = [
TypeBuilder<(ins CArg<"Register", "Register(0)">:$reg), [{
return $_get($_ctxt, normalizeRegister(reg));
}]>
];
let skipDefaultBuilders = 1;
string declarations = StrSubst<[{
/// The register kind for this SREG type.
static constexpr RegisterKind kRegisterKind = RegisterKind::$kind;
}], [VarRepl<"kind", kind>]>.result;
let extraClassDeclaration = declarations # [{
/// Returns true if the register is relocatable.
bool isRelocatable() const { return false; }
static Register normalizeRegister(Register reg) {
if (reg.getSemantics() == RegisterSemantics::Unallocated)
return Register(0);
return reg;
}

//===------------------------------------------------------------------===//
// RegisterTypeInterface
//===------------------------------------------------------------------===//
bool isRegisterRange() const { return false; }
RegisterRange getAsRange() const {
return RegisterRange(Register(0), 1);
return RegisterRange(getReg(), 1);
}
RegisterKind getRegisterKind() const {
return kRegisterKind;
}
RegisterTypeInterface cloneRegisterType(RegisterRange range) const {
return get(getContext());
}
RegisterTypeInterface cloneRegisterType(Register reg) const {
return get(getContext());
assert(range.size() == 1 && "SREG type can only clone single register");
return get(getContext(), range.begin());
}

//===------------------------------------------------------------------===//
Expand Down
41 changes: 8 additions & 33 deletions include/aster/Dialect/AMDGCN/Transforms/AMDGCNPasses.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -291,44 +291,19 @@ def ConvertLDSBuffers : Pass<"amdgcn-convert-lds-buffers"> {
];
}

def ConvertSCFControlFlow : Pass<"amdgcn-convert-scf-control-flow"> {
let summary = "Convert SCF control flow to AMDGCN control flow instructions";
let description = [{
This pass converts SCF structured control flow operations (such as scf.for,
scf.if, scf.while) to AMDGCN control flow instructions.

The pass first runs thread uniform analysis to determine whether loop
induction variables and conditions are uniform across all threads. Based on
this analysis:

- For thread-uniform conditions: emit scalar compare instructions (s_cmp_*)
and branch on SCC
- For thread-divergent conditions: emit vector compare instructions (v_cmp_*)
and branch on VCC/VCCZ

This pass should run after the ABI has been set and before register
allocation.

Post-condition: #amdgcn.no_scf_ops
}];
let dependentDialects = [
"mlir::aster::amdgcn::AMDGCNDialect",
"mlir::aster::lsir::LSIRDialect",
"mlir::cf::ControlFlowDialect"
];
}

def LegalizeCF : Pass<"amdgcn-legalize-cf"> {
let summary = "Legalize CF dialect ops to AMDGCN scalar branch instructions";
let description = [{
This pass legalizes CF dialect operations (cf.cond_br, cf.br) and lsir.cmpi
to AMDGCN scalar branch and compare instructions. It runs after register
allocation when operands are in physical registers and values flow through
side effects.
This pass legalizes CF dialect operations (cf.cond_br, cf.br) to AMDGCN
scalar branch instructions. It runs after register allocation when operands
are in physical registers and values flow through side effects.

The pass expects cf.cond_br conditions to come from amdgcn.is_cc which
tests an SCC or VCC register.

Transformations:
- lsir.cmpi (returns i1) -> s_cmp_* (sets SCC flag)
- cf.cond_br -> s_cbranch_scc1 / scc0 + s_branch
- cf.cond_br (cond from amdgcn.is_cc) -> s_cbranch_scc1/scc0 or
s_cbranch_vccnz/vccz + s_branch
- cf.br -> s_branch

Pre-condition: #amdgcn.all_registers_allocated
Expand Down
1 change: 1 addition & 0 deletions include/aster/Dialect/LSIR/IR/LSIROps.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dialect.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/InferIntRangeInterface.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
Expand Down
Loading