diff --git a/jaxlib/mosaic/dialect/gpu/mosaic_gpu.h b/jaxlib/mosaic/dialect/gpu/mosaic_gpu.h index ad1a6abd4..b4f13c50b 100644 --- a/jaxlib/mosaic/dialect/gpu/mosaic_gpu.h +++ b/jaxlib/mosaic/dialect/gpu/mosaic_gpu.h @@ -46,10 +46,6 @@ namespace mosaic_gpu { using Memref = ::mlir::TypedValue<::mlir::MemRefType>; using Pointer = ::mlir::TypedValue<::mlir::LLVM::LLVMPointerType>; -struct GlobalMemory : public mlir::SideEffects::Resource::Base { - llvm::StringRef getName() final { return ""; } -}; - constexpr absl::string_view kRuntimeTmaDescriptorInitializerName = "mosaic_gpu_init_tma_desc"; diff --git a/jaxlib/mosaic/dialect/gpu/mosaic_gpu.td b/jaxlib/mosaic/dialect/gpu/mosaic_gpu.td index fb02da12e..eb26c811b 100644 --- a/jaxlib/mosaic/dialect/gpu/mosaic_gpu.td +++ b/jaxlib/mosaic/dialect/gpu/mosaic_gpu.td @@ -243,10 +243,8 @@ def LayoutAttr : MosaicGPU_Attr<"Layout", "layout", let assemblyFormat = "`<` $num_dimensions `,` $transforms `>`"; } -def GlobalMemory : Resource<"::mosaic_gpu::GlobalMemory">; - def MosaicGPU_AsyncLoadOp : Op]>]> { + [AttrSizedOperandSegments]> { let summary = "Schedules an async load of a MemRef from GMEM to SMEM"; let description = [{ Schedules an async copy of the contents of the `source` MemRef in GMEM to @@ -313,7 +311,7 @@ def MosaicGPU_AsyncLoadOp : Op]>]> { + [AttrSizedOperandSegments]> { let summary = "Schedules an async store of a MemRef from SMEM to GMEM"; let description = [{ Schedules an async store of the contents of the `source` MemRef in SMEM to