[Wave] Add support for direct global load to lds

Signed-off-by: nithinsubbiah <nithinsubbiah@gmail.com>

Author: nithinsubbiah

iree/turbine/kernel/ops/wave_ops.py

@@ -287,6 +287,16 @@ def select(cond: "Register", if_true: "Register", if_false: "Register") -> "Regi
     ...
 
 
+def gather_to_lds(
+    src: "Memory",
+    src_idx: dict[IndexSymbol, IndexSequence],
+    dst: "Memory",
+    dst_idx: dict[IndexSymbol, IndexSequence],
+    dtype: DataType,
+):
+    ...
+
+
 def define_op(op_name: str) -> Callable[[T], T]:
     def decorator(cls: T) -> T:
         cls.tkw_op_name = op_name
@@ -2323,3 +2333,18 @@ def indexing_dims(self) -> list[IndexExpr]:
 
     def infer_type(self):
         self.type = get_custom(_to_sequence(self.args)[0]).type
+
+
+@define_op("gather_to_lds")
+@dataclass
+class GatherToLDS(CustomOp):
+    """
+    Represents an instruction that performs direct load from global
+    to lds.
+    """
+
+    src: Memory
+    src_idx: dict[IndexSymbol, IndexSequence]
+    dst: Memory
+    dst_idx: dict[IndexSymbol, IndexSequence]
+    dtype: DataType

iree/turbine/kernel/wave/codegen/handlers.py

@@ -60,6 +60,7 @@
     exp2,
     extract,
     extract_slice,
+    gather_to_lds,
     ge,
     get_custom,
     get_result,
@@ -1584,3 +1585,19 @@ def handle_reshape(emitter: WaveEmitter, node: fx.Node):
         [1],
     )
     emitter.bind_node_proxy(node, IRProxyValue(slice))
+
+
+@handle_op(gather_to_lds)
+def handle_gather_to_lds(emitter: WaveEmitter, node: fx.Node):
+    try:
+        src, src_idx, dst, dst_idx, dtype = node.args
+    except ValueError as e:
+        raise ValidationError("Malformed arguments") from e
+
+    n = amdgpu_d.gather_to_lds(
+        transfer_type=dtype, src=src, src_indices=src_idx, dst=dst, dst_indices=dst_idx
+    )
+
+    return n
+
+    # emitter.bind_node_proxy(node, IRProxyValue(n))

iree/turbine/kernel/wave/gather_to_shared.py

@@ -0,0 +1,63 @@
+# Copyright 2025 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from .._support.tracing import CapturedTrace
+from ..lang.global_symbols import *
+from ..ops.wave_ops import GatherToLDS, Write, get_custom
+from ..wave.constraints import (
+    Constraint,
+)
+from ..wave.utils.run_utils import get_default_arch
+from .utils.graph_utils import DCE, is_valid_global_read
+from .utils.symbol_utils import (
+    subs_idxc,
+)
+
+
+gather_to_shared_supported_arch = ["gfx950"]
+
+
+def get_write_node_info(read_node):
+    read_custom = get_custom(read_node)
+    write_node = None
+    for user in read_custom.users:
+        if (
+            isinstance(user, Write)
+            and subs_idxc(user.memory_type.address_space) == SHARED_ADDRESS_SPACE
+        ):
+            # get memory location and idx
+            dst = user.memory
+            dst_idx = user.get_derived_indices
+            write_node = user
+
+    return write_node, dst, dst_idx
+
+
+def gather_to_shared(trace: CapturedTrace, constraints: list[Constraint]):
+    """
+    This function enables direct memory load from global to lds without
+    passing through register reducing the data movement. This instruction
+    is supported only on specific architectures (gfx950).
+    """
+
+    # if get_default_arch() not in gather_to_shared_supported_arch:
+    #     return
+
+    global_read_nodes = trace.walk(is_valid_global_read)
+    for read_node in global_read_nodes:
+        custom = get_custom(read_node)
+        src = custom.memory
+        src_idx = custom.get_derived_indices
+        element_type = custom.type.dtype
+        write_node, dst, dst_idx = get_write_node_info(read_node)
+        with custom.graph.inserting_before(custom.fx_node):
+            write_node.replace_all_uses_with(
+                GatherToLDS(src, src_idx, dst, dst_idx, element_type).add_to_graph(
+                    custom.graph
+                )
+            )
+
+    DCE(trace)

iree/turbine/kernel/wave/global_to_shared_gathers.py

@@ -22,13 +22,15 @@
 from .utils.symbol_utils import subs_idxc
 from .utils.general_utils import is_gather
 from .minimize_global_loads import (
-    has_write_shared_user,
     construct_min_global_access_pattern,
     materialize_shape,
     identify_optimizable_loads,
     update_write_dependencies,
     SharedReadMetadata,
 )
+from .utils.graph_utils import (
+    has_write_shared_user,
+)
 
 """
 We are given N global gathers that are promoted to shared memory. This function

iree/turbine/kernel/wave/minimize_global_loads.py

@@ -11,7 +11,7 @@
     TilingConstraint,
 )
 from .._support.tracing import CapturedTrace
-from .._support.indexing import IndexingContext, IndexSequence, IndexSymbol, IndexExpr
+from .._support.indexing import IndexSequence, IndexSymbol, IndexExpr
 from ..ops.wave_ops import Read, Write, get_custom
 from ..lang.global_symbols import *
 from .utils.general_utils import (
@@ -20,6 +20,9 @@
     is_shared_read,
     get_fastest_index,
 )
+from .utils.graph_utils import (
+    is_valid_global_read,
+)
 from .utils.graph_utils import (
     DCE,
 )
@@ -41,23 +44,6 @@ class SharedReadMetadata:
     memory_shape: tuple[int | IndexExpr]
 
 
-def has_write_shared_user(node: Read) -> bool:
-    return any(
-        isinstance(user, Write)
-        and subs_idxc(user.memory_type.address_space) == SHARED_ADDRESS_SPACE
-        for user in node.users
-    )
-
-
-def is_valid_global_read(node: fx.Node) -> bool:
-    custom = get_custom(node)
-    return (
-        isinstance(custom, Read)
-        and subs_idxc(custom.memory_type.address_space) == GLOBAL_ADDRESS_SPACE
-        and has_write_shared_user(custom)
-    )
-
-
 def is_transposed_read(custom: Read) -> bool:
     """
     Checks whether or not we are doing a transposed read.

iree/turbine/kernel/wave/promotion.py

@@ -67,8 +67,8 @@ def apply_promotion_pattern(
     ```
     read_from_global lhs
     write_to_shared lhs
-    read_from_global lhs
-    write_to_shared lhs
+    read_from_global rhs
+    write_to_shared rhs
     shared_barrier
     read_from_shared lhs
     read_from_shared rhs

iree/turbine/kernel/wave/utils/graph_utils.py

@@ -9,6 +9,7 @@
 import iree.turbine.kernel.lang as tkl
 from ...ops.wave_ops import (
     get_custom,
+    Read,
     Write,
     NestedRegionOp,
     Output,
@@ -438,3 +439,23 @@ def is_barrier_between(src: fx.Node, dst: fx.Node) -> bool:
             return dst_check or root_check
 
         assert False, "Unhandled case when src and dst are in different graphs"
+
+
+def has_write_shared_user(node: Read) -> bool:
+    return any(
+        isinstance(user, Write)
+        and subs_idxc(user.memory_type.address_space) == SHARED_ADDRESS_SPACE
+        for user in node.users
+    )
+
+
+def is_valid_global_read(node: fx.Node) -> bool:
+    """
+    Check if a read node is global and if its user writes to shared memory.
+    """
+    custom = get_custom(node)
+    return (
+        isinstance(custom, Read)
+        and subs_idxc(custom.memory_type.address_space) == GLOBAL_ADDRESS_SPACE
+        and has_write_shared_user(custom)
+    )

iree/turbine/kernel/wave/utils/run_utils.py

@@ -236,4 +236,5 @@ def set_default_run_config(options: WaveCompileOptions) -> WaveCompileOptions:
     options.backend = "rocm"
     options.device = "hip"
     options.target = get_default_arch()
+
     return options

iree/turbine/kernel/wave/wave.py

@@ -12,7 +12,7 @@
 from ..lang import Grid, IndexMapping
 from ..lang.global_symbols import *
 from ..ops import wave_ops
-from ..ops.wave_ops import Iterate, CustomOp, get_custom, IterArg
+from ..ops.wave_ops import Iterate, CustomOp, get_custom
 from .._support.indexing import IndexingContext, IndexExpr
 from .symbolic_constraints import SymbolicAlias
 from .._support.tracing import (
@@ -50,6 +50,7 @@
 from .decompose_scan_ops import decompose_scan_ops
 from .decompose_dot_mma import decompose_dot_mma
 from .expansion.expansion import expand_graph, add_get_results
+from .gather_to_shared import gather_to_shared
 from .global_to_shared_gathers import global_to_shared_gathers
 from .hoisting import hoist_loop_invariant_ops
 from .minimize_global_loads import minimize_global_loads
@@ -86,12 +87,6 @@
 import inspect
 import sympy
 import warnings
-from pathlib import Path
-import sys
-import subprocess
-import os
-import shutil
-import glob
 
 __all__ = ["wave", "wave_trace_only"]
 
@@ -529,14 +524,15 @@ def _trace_and_get_kernel_signature(
         graph_passes = self.build_initial_pass_pipeline(
             trace, options, print_ir_before, print_ir_after
         )
-
+        breakpoint()
         # Optimizations.
         graph_passes += [
             partial(decompose_vmma_ops, trace, self.constraints),
             partial(decompose_dot_mma, trace, self.constraints),
             partial(hoist_loop_invariant_ops, trace, self.constraints),
             partial(global_to_shared_gathers, trace, self.constraints),
             partial(minimize_global_loads, trace, self.constraints),
+            partial(gather_to_shared, trace, self.constraints),
             partial(apply_shared_memory_indexing_corrections, trace, self.constraints),
         ]