Skip to content

Users/nico/moe step 0 permute bug #722

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
nicolasvasilache wants to merge 18 commits into
base: main
Choose a base branch
from
Draft

Users/nico/moe step 0 permute bug #722

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
2a5a076
[Wave] Add a GEMM-bias add example
harsh-nod Apr 2, 2025
79cb1f0
update partition strided operators to handle read
harsh-nod Apr 3, 2025
f89b69a
Add MOE kernel - Step 1
nicolasvasilache Mar 4, 2025
dd6f32a
Even smaller step where we only compute the last step of the MoE.
nicolasvasilache Apr 1, 2025
524e898
Minimal changes to allow vector<2xt> load/stores
nicolasvasilache Apr 7, 2025
6304d39
Undo load-bearing resolve_thread_shapes changes
nicolasvasilache Apr 11, 2025
392c94f
Valid but incorrect IR.
nicolasvasilache Apr 11, 2025
ed7a154
Drop manually specified elements_per_thread
nicolasvasilache Apr 11, 2025
03d41ca
[NFC] Factor out pass creation in wave.py
nicolasvasilache Apr 15, 2025
67fa3a5
Use factored out pass creation in moe.py test and uniformize IR creat…
nicolasvasilache Apr 15, 2025
bbd32e7
Step debugging using the system as a black box.
nicolasvasilache Apr 16, 2025
db14b63
[NFC] Simplify partition_strided_operators.py
nicolasvasilache Apr 16, 2025
715a642
Fix offset_mapping_w2 IndexMapping
nicolasvasilache Apr 17, 2025
0a1c64b
[NFC] Take reshape out of torch_naive_moe_step_3 and simplify API
nicolasvasilache Apr 17, 2025
c79f3f9
Set TOPK: 0 in vector_shapes to specify a batch dimension
nicolasvasilache Apr 23, 2025
5e5d91a
[NFC] tkw.reduction -> tkw.iterate post rebase
nicolasvasilache Apr 23, 2025
4cedbe8
Hack in batch dims in tkw.mma
nicolasvasilache Apr 25, 2025
df7f4b9
Repro permutation issue
nicolasvasilache Apr 17, 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
137 changes: 58 additions & 79 deletions iree/turbine/kernel/ops/wave_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,65 +42,57 @@

def allocate(
shape: tuple[IndexExpr], dtype: DataType, address_space: IndexSymbol
) -> "Memory":
...
) -> "Memory": ...


def self_index(
dim: IndexExpr,
dtype: DataType,
elements_per_thread: Optional[IndexExpr | int] = None,
) -> "Register":
...
) -> "Register": ...


def extract(
register: "Register",
offsets: tuple[IndexExpr],
) -> "Register":
...
) -> "Register": ...


def extract_slice(
register: "Register",
offsets: tuple[IndexExpr],
sizes: tuple[IndexExpr],
strides: tuple[IndexExpr],
) -> "Register":
...
) -> "Register": ...


def shared_memory_barrier():
...
def shared_memory_barrier(): ...


def read(
memory: "Memory",
elements_per_thread: Optional[IndexExpr | int] = None,
mapping: Optional[IndexMapping] = None,
mapping_dynamic_vals: "Register" | tuple["Register", ...] = (),
) -> "Register":
...
) -> "Register": ...


def conditional(
condition: "Register" | IndexExpr,
) -> Callable[[Callable[[], None]], None]:
...
) -> Callable[[Callable[[], None]], None]: ...


def iterate(
axis: IndexExpr, init_args: Sequence["Register"]
) -> Callable[[Callable[[AccT], AccT]], AccT]:
...
) -> Callable[[Callable[[AccT], AccT]], AccT]: ...


def register(shape: tuple[IndexExpr, ...], dtype: DataType, value: float) -> "Register":
...
def register(
shape: tuple[IndexExpr, ...], dtype: DataType, value: float
) -> "Register": ...


def mma(lhs: "Register", rhs: "Register", acc: "Register") -> "Register":
...
def mma(lhs: "Register", rhs: "Register", acc: "Register") -> "Register": ...


def write(
Expand All @@ -109,32 +101,27 @@ def write(
elements_per_thread: Optional[IndexExpr | int] = None,
mapping: Optional[IndexMapping] = None,
mapping_dynamic_vals: "Register" | tuple["Register", ...] = (),
):
...
): ...


def apply_expr(value: "Register" | Sequence["Register"], expr: Callable) -> "Register":
...
def apply_expr(
value: "Register" | Sequence["Register"], expr: Callable
) -> "Register": ...


def set_symbol(symbol: IndexExpr, value: "Register"):
...
def set_symbol(symbol: IndexExpr, value: "Register"): ...


def exp2(src: "Register") -> "Register":
...
def exp2(src: "Register") -> "Register": ...


def log2(src: "Register") -> "Register":
...
def log2(src: "Register") -> "Register": ...


def reciprocal(src: "Register") -> "Register":
...
def reciprocal(src: "Register") -> "Register": ...


def abs(src: "Register") -> "Register":
...
def abs(src: "Register") -> "Register": ...


def tanh_approx(src: "Register") -> "Register":
Expand All @@ -145,88 +132,73 @@ def tanh(src: "Register") -> "Register":
...


def roundeven(src: "Register") -> "Register":
...
def roundeven(src: "Register") -> "Register": ...


def maximum(lhs: "Register", rhs: "Register") -> "Register":
...
def maximum(lhs: "Register", rhs: "Register") -> "Register": ...


def minimum(lhs: "Register", rhs: "Register") -> "Register":
...
def minimum(lhs: "Register", rhs: "Register") -> "Register": ...


def broadcast(
arg: "Register", target_shape: Optional[Sequence[IndexExpr | int]] = None
) -> "Register":
...
) -> "Register": ...


def sum(
src: "Register",
acc: Optional["Register"] = None,
dim: Optional[IndexExpr | int] = None,
) -> "Register":
...
) -> "Register": ...


def max(
src: "Register",
acc: Optional["Register"] = None,
dim: Optional[IndexExpr | int] = None,
) -> "Register":
...
) -> "Register": ...


def min(
src: "Register",
acc: Optional["Register"] = None,
dim: Optional[IndexExpr | int] = None,
) -> "Register":
...
) -> "Register": ...


def shuffle(src: "Register", offset: int, width: int) -> "Register":
...
def shuffle(src: "Register", offset: int, width: int) -> "Register": ...


def gt(lhs: "Register", rhs: "Register") -> "Register":
...
def gt(lhs: "Register", rhs: "Register") -> "Register": ...


def ge(lhs: "Register", rhs: "Register") -> "Register":
...
def ge(lhs: "Register", rhs: "Register") -> "Register": ...


def lt(lhs: "Register", rhs: "Register") -> "Register":
...
def lt(lhs: "Register", rhs: "Register") -> "Register": ...


def le(lhs: "Register", rhs: "Register") -> "Register":
...
def le(lhs: "Register", rhs: "Register") -> "Register": ...


def eq(lhs: "Register", rhs: "Register") -> "Register":
...
def eq(lhs: "Register", rhs: "Register") -> "Register": ...


def cast(src: "Register", dtype: DataType) -> "Register":
...
def cast(src: "Register", dtype: DataType) -> "Register": ...


def permute(src: "Register", target_shape: Sequence[IndexExpr]) -> "Register":
...
def permute(src: "Register", target_shape: Sequence[IndexExpr]) -> "Register": ...


def reshape(
inputs: Sequence["Register"], target_vector_shape: dict[IndexSymbol, int]
) -> "Register":
...
) -> "Register": ...


def select(cond: "Register", if_true: "Register", if_false: "Register") -> "Register":
...
def select(
cond: "Register", if_true: "Register", if_false: "Register"
) -> "Register": ...


def define_op(op_name: str) -> Callable[[T], T]:
Expand Down Expand Up @@ -430,7 +402,10 @@ def custom_string(self, value_map: dict[str, str]) -> str:
if hasattr(self.fx_node, "index") and self.fx_node.index:
vars_list.append(f"index={self.fx_node.index}")
vars_str = ", ".join(vars_list)
return f"{self.tkw_op_name}({vars_str}) type({self.fx_node.type})"
return f"""{self.tkw_op_name}({vars_str})
type({self.fx_node.type})
indexing_dims({self.indexing_dims if hasattr(self, 'indexing_dims') else None})
vector_shapes({self.vector_shapes if hasattr(self, 'vector_shapes') else None})"""

def add_to_graph(self, region_graph: RegionGraph, type: Any = None) -> fx.Node:
arg_list = tuple([value for _, value in vars(self).items()])
Expand Down Expand Up @@ -1226,7 +1201,10 @@ def custom_string(self, value_map: dict[str, str]) -> str:
custom_str += f"lhs={self.lhs} (index = {self.lhs_index}), "
custom_str += f"rhs={self.rhs} (index = {self.rhs_index}), "
custom_str += f"acc={self.acc} (index = {self.acc_index}))"
custom_str += f" type({self.fx_node.type})"
custom_str += f"\n\ttype({self.fx_node.type})"
custom_str += f"\n\tindex({self.index})"
custom_str += f"\n\tindexing_dims({self.indexing_dims if hasattr(self, 'indexing_dims') else None})"
custom_str += f"\n\tvector_shapes({self.vector_shapes if hasattr(self, 'vector_shapes') else None})"
return custom_str

def align_index(self, constraints: list["Constraint"]) -> None:
Expand Down Expand Up @@ -1529,20 +1507,19 @@ def outputs(self, graph: fx.Graph) -> list[fx.Node]:

@property
def index(self) -> list[dict[IndexSymbol, IndexSequence]]:
def get_index(node: fx.Node) -> dict[IndexSymbol, IndexSequence]:
custom = get_custom(node)
if isinstance(custom, MMA):
return custom.acc_index
return custom.index

for node in self.get_root_graph().subgraphs[self.subgraph_name].nodes:
if isinstance(output := get_custom(node), Output):
return_vals = output.return_vals[0]
return (
[
(
get_custom(val).acc_index
if isinstance(get_custom(val), MMA)
else val.index
)
for val in return_vals
]
[get_index(val) for val in return_vals]
if isinstance(return_vals, (Sequence))
else return_vals.index
else get_index(return_vals)
)

@index.setter
Expand Down Expand Up @@ -1759,6 +1736,8 @@ def index(self) -> dict[IndexSymbol, IndexSequence]:
return None
if not isinstance(custom, Iterate):
return custom_index
if not isinstance(custom_index, Sequence):
return custom_index
assert isinstance(custom_index, Sequence) and self.res_idx < len(
custom.indexing_dims
), f"Invalid {custom_index=} with {self.res_idx=} and {custom.indexing_dims=}\n{custom}"
Expand Down
10 changes: 7 additions & 3 deletions iree/turbine/kernel/wave/analysis/index_sequence_analysis.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
BinaryPyOp,
Broadcast,
CustomOp,
GetResult,
IterArg,
MMA,
NestedRegionOp,
Expand Down Expand Up @@ -37,9 +38,7 @@
get_workgroup_constraints,
partial,
)
from ..utils.mma_utils import (
get_mma_dimensional_mapping,
)
from ..utils.mma_utils import get_mma_dimensional_mapping
from ..utils.graph_utils import (
get_inputs,
get_users,
Expand Down Expand Up @@ -462,6 +461,11 @@ def propagate_indices(
visited: set[CustomOp],
symbolic_constraints: list[SymbolicAlias],
):
def get_index(custom: CustomOp):
if isinstance(custom, MMA):
return custom.acc.index
return custom.index

"""
Propagate the index and vector shapes through the graph
starting with priveleged nodes (like MMA, Read, Write).
Expand Down
Loading