[lang] Remove nvvm_mlir_interfaces.py

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

.flake8

@@ -16,7 +16,6 @@ exclude =
     *docs/venv,
     *docs/source/_templates,
     experimental/cuda-lang/src/cuda/lang/_stub/nvvm.py,
-    experimental/cuda-lang/src/cuda/lang/_stub/nvvm_mlir_interfaces.py,
     experimental/cuda-lang/src/cuda/lang/_stub/_libdevice.py,
 per-file-ignores =
     src/cuda/tile/_stub.py: W291

experimental/cuda-lang/src/cuda/lang/_ir/op_impl/core_api_impl.py

@@ -5,6 +5,7 @@
 from cuda.lang._enums import CachePolicy
 from cuda.lang._exception import InvalidValueError, TypeCheckingError
 from cuda.lang._ir.op_defs import RawNVVMIntrinsic, BitCast, InlinePTX
+from cuda.tile import MemoryScope
 from ..type import (
     DTypeConstructor,
     MemorySpace,
@@ -322,3 +323,34 @@ def impl_create_fractional_cache_policy(
         read_write_operands=(),
     )
     return results[0]
+
+
+@impl(core_api.memory_barrier)
+def memory_barrier_impl(scope: Var) -> None:
+    scope2intrin = {
+        MemoryScope.BLOCK: "llvm.nvvm.membar.cta",
+        MemoryScope.CLUSTER: "llvm.nvvm.fence_sc_cluster",
+        MemoryScope.DEVICE: "llvm.nvvm.membar.gl",
+        MemoryScope.SYS: "llvm.nvvm.membar.sys",
+    }
+    scope = require_constant_enum(scope, MemoryScope)
+    intrinsic = scope2intrin.get(scope)
+    if intrinsic is None:
+        valid = ", ".join(x._name_ for x in scope2intrin.keys())
+        raise InvalidValueError(f"Invalid memory scope '{scope}'. Valid values: {valid}")
+    add_operation_variadic(
+        RawNVVMIntrinsic, (),
+        intrinsic=intrinsic,
+        operands_=()
+    )
+
+
+@impl(core_api.grid_dependency_control_wait, fixed_args=["wait"])
+@impl(core_api.grid_dependency_control_launch_dependents, fixed_args=["launch.dependents"])
+def grid_dependency_control_action_impl(action: str) -> None:
+    add_operation_variadic(
+        RawNVVMIntrinsic, (),
+        intrinsic="llvm.nvvm."
+                  "griddepcontrol." + action,
+        operands_=()
+    )

experimental/cuda-lang/src/cuda/lang/_ir/op_impl/prefetch_impl.py

@@ -0,0 +1,104 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+from cuda.lang._enums import PrefetchLevel, CachePolicy
+from cuda.lang._ir.op_defs import RawNVVMIntrinsic
+from cuda.lang._ir.type import TensorMapTy
+from cuda.lang._ir.type_checking_helpers import require_pointer_type
+from cuda.lang._stub.prefetch import prefetch, prefetch_uniform, prefetch_tensor_map
+from cuda.tile._datatype import PointerInfo
+from cuda.tile._exception import InvalidValueError, TypeCheckingError
+from cuda.tile._ir.ir import Var, add_operation_variadic
+from cuda.tile._ir.op_impl import ImplRegistry, require_constant_enum
+from cuda.tile._ir.type import NONE
+from cuda.tile._memory_model import MemorySpace
+
+_registry = ImplRegistry()
+impl = _registry.impl
+
+
+def prefetch_impl_registry() -> ImplRegistry:
+    return _registry
+
+
+@impl(prefetch)
+def prefetch_impl(address: Var, level: Var, eviction_priority: Var):
+    address_ty = require_pointer_type(address)
+    address_space = PointerInfo(address_ty.pointer_dtype).memory_space
+    level = require_constant_enum(level, PrefetchLevel)
+
+    if eviction_priority.get_type() == NONE:
+        evict_str = ""
+    else:
+        eviction_priority = require_constant_enum(eviction_priority, CachePolicy)
+        if level != PrefetchLevel.L2:
+            raise InvalidValueError("Prefetch eviction priority is supported only for L2")
+        priority_map = {CachePolicy.L2_EVICT_NORMAL: "normal", CachePolicy.L2_EVICT_LAST: "last"}
+        suffix = priority_map.get(eviction_priority)
+        if suffix is None:
+            valid = ", ".join(x._name_ for x in priority_map.keys())
+            raise InvalidValueError(f"Invalid eviction priority {eviction_priority._name_}."
+                                    f" Accepted values: {valid}.")
+        if address_space != MemorySpace.GLOBAL:
+            raise TypeCheckingError(f"Pointer in GLOBAL address space is required"
+                                    f" when eviction priority is specified;"
+                                    f" received a {address_space._name_} pointer instead.")
+        evict_str = f".evict.{suffix}"
+
+    address_space_map = {
+        MemorySpace.GENERIC: "",
+        MemorySpace.GLOBAL: ".global",
+        MemorySpace.LOCAL: ".local",
+    }
+    address_space_str = address_space_map.get(address_space)
+    if address_space_str is None:
+        valid = ", ".join(x._name_ for x in address_space_map.keys())
+        raise TypeCheckingError(f"Invalid pointer address space {address_space._name_}."
+                                f" Accepted address spaces: {valid}.")
+
+    level_map = {PrefetchLevel.L1: "L1", PrefetchLevel.L2: "L2"}
+    level_str = level_map[level]
+
+    add_operation_variadic(
+        RawNVVMIntrinsic,
+        (),
+        intrinsic=f"llvm.nvvm.prefetch{address_space_str}.{level_str}{evict_str}",
+        operands_=(address,)
+    )
+
+
+@impl(prefetch_uniform)
+def prefetch_uniform_impl(address: Var):
+    address_ty = require_pointer_type(address)
+    address_space = PointerInfo(address_ty.pointer_dtype).memory_space
+    if address_space != MemorySpace.GENERIC:
+        raise TypeCheckingError(f"Expected a pointer in GENERIC address space;"
+                                f" received {address_space._name_} instead.")
+
+    add_operation_variadic(
+        RawNVVMIntrinsic,
+        (),
+        intrinsic="llvm.nvvm.prefetchu.L1",
+        operands_=(address,)
+    )
+
+
+@impl(prefetch_tensor_map)
+def prefetch_tensor_map_impl(tensor_map: Var):
+    if isinstance(tensor_map.get_type(), TensorMapTy):
+        from cuda.lang._ir.ops import tensor_map_as_opaque_ptr
+        tensor_map = tensor_map_as_opaque_ptr(tensor_map)
+    else:
+        address_ty = require_pointer_type(tensor_map)
+        address_space = PointerInfo(address_ty.pointer_dtype).memory_space
+        valid_spaces = (MemorySpace.GENERIC, MemorySpace.CONSTANT)
+        if address_space not in valid_spaces:
+            valid = ", ".join(x._name_ for x in valid_spaces)
+            raise TypeCheckingError(f"Invalid address space {address_space._name_}."
+                                    f" Accepted address spaces: {valid}")
+    add_operation_variadic(
+        RawNVVMIntrinsic,
+        (),
+        intrinsic="llvm.nvvm.prefetch.tensormap",
+        operands_=(tensor_map,)
+    )

experimental/cuda-lang/src/cuda/lang/_ir/ops.py

@@ -165,6 +165,7 @@
     contiguous_strides_from_shape,
     pointer_impl_registry,
 )
+from .op_impl.prefetch_impl import prefetch_impl_registry
 from .op_impl.copy_async_impl import copy_async_impl_registry
 from .op_impl.barrier_impl import barrier_impl_registry
 from .op_impl.mbarrier_impl import mbarrier_impl_registry
@@ -187,6 +188,7 @@
 cuda_lang_impl_registry.update(fence_impl_registry())
 cuda_lang_impl_registry.update(vector_impl_registry())
 cuda_lang_impl_registry.update(pointer_impl_registry())
+cuda_lang_impl_registry.update(prefetch_impl_registry())
 cuda_lang_impl_registry.update(copy_async_impl_registry())
 cuda_lang_impl_registry.update(barrier_impl_registry())
 cuda_lang_impl_registry.update(mbarrier_impl_registry())
@@ -786,11 +788,15 @@ def tensor_map_get_transaction_bytes_impl(self: Var, mode: Var):
     return loosely_typed_const(element_count * map_ty.element_bitwidth // 8)
 
 
+def tensor_map_as_opaque_ptr(self: Var):
+    result_ty = PointerTy(opaque_pointer_dtype())
+    return add_operation(TensorMapAsOpaquePtr, result_ty, tensor_map=self)
+
+
 @impl(tensor_map.TensorMap.as_opaque_ptr)
 def tensor_map_as_opaque_ptr_impl(self: Var):
     require_tensor_map_ty(self)
-    result_ty = PointerTy(opaque_pointer_dtype())
-    return add_operation(TensorMapAsOpaquePtr, result_ty, tensor_map=self)
+    return tensor_map_as_opaque_ptr(self)
 
 
 @impl(cluster_launch_control_try_cancel)

experimental/cuda-lang/src/cuda/lang/_stub/core_api.py

@@ -1124,23 +1124,19 @@ def nanosleep(nanoseconds: int):
     nvvm.nanosleep(nanoseconds)
 
 
+@stub
 def memory_barrier(scope: MemoryScope) -> None:
     """Issue a memory fence with the given scope."""
-    return nvvm_mlir_interfaces.memory_barrier(scope=scope)
 
 
+@stub
 def grid_dependency_control_wait() -> None:
     """Wait for prerequisite grids in a programmatic dependent launch."""
-    nvvm_mlir_interfaces.griddepcontrol(
-        kind=nvvm_mlir_interfaces.GridDepActionKind.wait
-    )
 
 
+@stub
 def grid_dependency_control_launch_dependents() -> None:
     """Launch dependent grids in a programmatic dependent launch."""
-    nvvm_mlir_interfaces.griddepcontrol(
-        kind=nvvm_mlir_interfaces.GridDepActionKind.launch_dependents
-    )
 
 
 @stub
@@ -1189,4 +1185,3 @@ def cdiv(x, y, /):
 
 # Need these imports at the end in order to overcome the circular import problem
 from . import nvvm  # noqa: E402
-from . import nvvm_mlir_interfaces  # noqa: E402