[Mosaic GPU] Delete unused declarations of `mosaic_gpu_memcpy_async_h…

…2d`. PiperOrigin-RevId: 716616807
jax-ml · Jan 17, 2025 · d3be190 · d3be190
1 parent d34c40f
commit d3be190
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 16 deletions.
diff --git a/jax/experimental/mosaic/gpu/core.py b/jax/experimental/mosaic/gpu/core.py
@@ -488,10 +488,6 @@ def _declare_runtime_functions():
   func.FuncOp(
       "mosaic_gpu_init_tma_desc", init_tma_desc_type, visibility="private"
   )
-  memcpy_async_type = ir.FunctionType.get([ptr_ty, ptr_ty, i64, ptr_ty], [])
-  func.FuncOp(
-      "mosaic_gpu_memcpy_async_h2d", memcpy_async_type, visibility="private"
-  )
 
 
 def as_gpu_kernel(

diff --git a/jaxlib/mosaic/dialect/gpu/mosaic_gpu.cc b/jaxlib/mosaic/dialect/gpu/mosaic_gpu.cc
@@ -231,12 +231,6 @@ void DeclareRuntimeFunctions(mlir::OpBuilder& builder) {
           builder.getFunctionType(
               TypeRange{ptr, ptr, i64, i64, ptr, ptr, i64, ptr}, TypeRange{}))
       .setVisibility(mlir::func::FuncOp::Visibility::Private);
-
-  builder
-      .create<mlir::func::FuncOp>(
-          builder.getUnknownLoc(), kRuntimeMemcpyAsyncH2DName,
-          builder.getFunctionType(TypeRange{ptr, ptr, i64, ptr}, TypeRange{}))
-      .setVisibility(mlir::func::FuncOp::Visibility::Private);
 }
 
 bool IsContiguous(mlir::MemRefType type) {

diff --git a/jaxlib/mosaic/dialect/gpu/mosaic_gpu.h b/jaxlib/mosaic/dialect/gpu/mosaic_gpu.h
@@ -52,8 +52,6 @@ struct GlobalMemory : public mlir::SideEffects::Resource::Base<GlobalMemory> {
 
 constexpr absl::string_view kRuntimeTmaDescriptorInitializerName =
     "mosaic_gpu_init_tma_desc";
-constexpr absl::string_view kRuntimeMemcpyAsyncH2DName =
-    "mosaic_gpu_memcpy_async_h2d";
 
 template <typename T>
 std::string MlirToString(T&& value) {

diff --git a/jaxlib/mosaic/dialect/gpu/mosaic_gpu_test.cc b/jaxlib/mosaic/dialect/gpu/mosaic_gpu_test.cc
@@ -181,16 +181,16 @@ TEST_F(MosaicGpuTest, RuntimeFunctionsAreRegistered) {
 
   llvm::SmallVector<mlir::func::FuncOp> func_ops =
       llvm::to_vector(module_op->getBody()->getOps<mlir::func::FuncOp>());
-  EXPECT_EQ(func_ops.size(), 2);
+  EXPECT_EQ(func_ops.size(), 1);
 
   absl::flat_hash_set<std::string> func_names;
   for (mlir::func::FuncOp& func_op : func_ops) {
     func_names.insert(func_op.getSymName().str());
   }
 
-  EXPECT_THAT(func_names, UnorderedElementsAre(
-                              mosaic_gpu::kRuntimeTmaDescriptorInitializerName,
-                              mosaic_gpu::kRuntimeMemcpyAsyncH2DName));
+  EXPECT_THAT(
+      func_names,
+      UnorderedElementsAre(mosaic_gpu::kRuntimeTmaDescriptorInitializerName));
 }