[LLVMCPU] Lower explicit workgroup-local allocs through dispatch ABI

Add #iree_codegen.workgroup_local as a codegen memory-space attribute for authored LLVMCPU dispatch allocations that should use HAL workgroup local memory instead of the thread stack.

This deliberately matches the GPU-side semantic split: workgroup/shared memory is represented as memref.alloc in a workgroup memory space, while memref.alloca remains private stack scratch. LLVMCPU uses an IREE-specific memory-space attribute instead of #gpu.address_space<workgroup>, but keeps the same alloc/dealloc ownership model.

The assignment pass is intentionally narrow: it only handles entry-block memref.alloc operations in HAL executable exports, computes byte ranges with the target data layout and ABI alignment, rejects dynamic or unsupported layouts, and refuses to overwrite existing local-memory assignments or predeclared export requirements.

ConvertToLLVM consumes the assigned range by building memref descriptors from the HAL workgroup local-memory pointer. Matching memref.dealloc operations are erased because the storage is owned by the dispatch frame.

Author: benvanik

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h

@@ -73,6 +73,8 @@ constexpr StringLiteral kUkernelAttrName = "iree_codegen.ukernel";
 constexpr StringLiteral kUKernelProviderName = "iree_codegen.ukernel_provider";
 constexpr StringLiteral kVectorTileSizesAttrName =
     "iree_codegen.vector_tile_sizes";
+constexpr StringLiteral kWorkgroupLocalMemoryRangeAttrName =
+    "iree_codegen.local_memory_range";
 
 //===----------------------------------------------------------------------===//
 // Helpers for getting/setting iree_codegen.translation_info attribute on a

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td

@@ -160,6 +160,35 @@ def LocalMappingAttr :
   }];
 }
 
+//===---------------------------------------------------------------------===//
+// iree_codegen.workgroup_local memory space attribute
+//===---------------------------------------------------------------------===//
+
+def WorkgroupLocalMemoryAttr :
+    AttrDef<IREECodegen_Dialect, "WorkgroupLocalMemory"> {
+  let mnemonic = "workgroup_local";
+  let parameters = (ins);
+  let assemblyFormat = "";
+
+  let description = [{
+    Memory space attribute indicating that an allocation uses per-workgroup
+    local memory supplied by the target runtime instead of the thread stack.
+
+    LLVMCPU dispatch lowering rewrites `memref.alloc` operations marked with
+    this memory space to use the HAL dispatch ABI local-memory pointer. Matching
+    `memref.dealloc` operations are erased because the storage is owned by the
+    dispatch frame. The allocation layout is assigned before final LLVM
+    conversion and exported as the dispatch `workgroup_local_memory`
+    requirement.
+
+    Example:
+    ```mlir
+      %scratch = memref.alloc()
+        : memref<64x1536xf32, #iree_codegen.workgroup_local>
+    ```
+  }];
+}
+
 //===---------------------------------------------------------------------===//
 // iree_codegen.simple_target
 //===---------------------------------------------------------------------===//

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/test/roundtrip.mlir

@@ -114,6 +114,17 @@ func.func private @workgroup_scope_attr_linearize() attributes {
 
 // -----
 
+func.func @workgroup_local_memory_space(
+    %arg0: memref<64x64xf32, #iree_codegen.workgroup_local>)
+    -> memref<64x64xf32, #iree_codegen.workgroup_local> {
+  return %arg0 : memref<64x64xf32, #iree_codegen.workgroup_local>
+}
+// CHECK-LABEL: func.func @workgroup_local_memory_space(
+// CHECK-SAME:    memref<64x64xf32, #iree_codegen.workgroup_local>
+// CHECK-SAME:    -> memref<64x64xf32, #iree_codegen.workgroup_local>
+
+// -----
+
 // Test constraints op with knobs and dims.
 func.func @constraints_op(%arg0: index, %arg1: index) {
   iree_codegen.smt.constraints target = <set = 0>, pipeline = #iree_gpu.pipeline<VectorDistribute>,

compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel

@@ -60,6 +60,7 @@ iree_compiler_cc_library(
         "LLVMCPU2DScalableTo1DScalable.cpp",
         "LLVMCPUAssignConstantOrdinals.cpp",
         "LLVMCPUAssignImportOrdinals.cpp",
+        "LLVMCPUAssignWorkgroupLocalMemory.cpp",
         "LLVMCPUCheckIRBeforeLLVMConversion.cpp",
         "LLVMCPUEmitVectorizationRemarks.cpp",
         "LLVMCPULinkExecutables.cpp",

compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt

@@ -56,6 +56,7 @@ iree_cc_library(
     "LLVMCPU2DScalableTo1DScalable.cpp"
     "LLVMCPUAssignConstantOrdinals.cpp"
     "LLVMCPUAssignImportOrdinals.cpp"
+    "LLVMCPUAssignWorkgroupLocalMemory.cpp"
     "LLVMCPUCheckIRBeforeLLVMConversion.cpp"
     "LLVMCPUEmitVectorizationRemarks.cpp"
     "LLVMCPULinkExecutables.cpp"

compiler/src/iree/compiler/Codegen/LLVMCPU/ConvertToLLVM.cpp

@@ -6,6 +6,7 @@
 
 #include "iree/compiler/Codegen/Common/PassUtils.h"
 #include "iree/compiler/Codegen/Common/Transforms.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
 #include "iree/compiler/Codegen/LLVMCPU/DispatchABI.h"
 #include "iree/compiler/Codegen/LLVMCPU/Passes.h"
@@ -332,6 +333,83 @@ struct ConvertHALInterfaceBindingSubspanOp
   }
 };
 
+/// Rewrites memref.alloc with #iree_codegen.workgroup_local memory space to a
+/// memref descriptor backed by HAL dispatch workgroup local memory.
+struct ConvertWorkgroupLocalAllocOp
+    : ConvertOpToLLVMWithABIPattern<memref::AllocOp> {
+  ConvertWorkgroupLocalAllocOp(HALDispatchABI &abi,
+                               LLVMTypeConverter &typeConverter)
+      : ConvertOpToLLVMWithABIPattern(abi, typeConverter, /*benefit=*/2) {}
+
+  LogicalResult
+  matchAndRewrite(memref::AllocOp allocOp, memref::AllocOpAdaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    MemRefType memRefType = allocOp.getType();
+    if (!isa_and_nonnull<IREE::Codegen::WorkgroupLocalMemoryAttr>(
+            memRefType.getMemorySpace())) {
+      return failure();
+    }
+    if (!memRefType.hasStaticShape()) {
+      return allocOp.emitOpError(
+          "workgroup local memory allocations must have static shape");
+    }
+    SmallVector<int64_t> strides;
+    int64_t offset = 0;
+    if (failed(memRefType.getStridesAndOffset(strides, offset)) ||
+        ShapedType::isDynamic(offset) ||
+        llvm::any_of(strides, ShapedType::isDynamic)) {
+      return allocOp.emitOpError(
+          "workgroup local memory allocations must have static layout");
+    }
+
+    auto rangeAttr = allocOp->getAttrOfType<DenseI64ArrayAttr>(
+        kWorkgroupLocalMemoryRangeAttrName);
+    if (!rangeAttr || rangeAttr.size() != 2 || rangeAttr[0] < 0 ||
+        rangeAttr[1] < 0) {
+      return allocOp.emitOpError(
+          "missing valid iree_codegen.local_memory_range annotation");
+    }
+
+    Location loc = allocOp.getLoc();
+    Value basePtr = abi.loadWorkgroupLocalMemoryPtr(allocOp, rewriter);
+    Value byteOffset = LLVM::ConstantOp::create(
+        rewriter, loc, rewriter.getI64Type(), rangeAttr[0]);
+    Value offsetPtr = LLVM::GEPOp::create(
+        rewriter, loc, basePtr.getType(), rewriter.getI8Type(), basePtr,
+        byteOffset, LLVM::GEPNoWrapFlags::inbounds);
+
+    MemRefType strippedType =
+        MemRefType::get(memRefType.getShape(), memRefType.getElementType(),
+                        memRefType.getLayout());
+    auto desc = MemRefDescriptor::fromStaticShape(
+        rewriter, loc, *getTypeConverter(), strippedType, offsetPtr);
+    rewriter.replaceOp(allocOp, {desc});
+    return success();
+  }
+};
+
+/// Erases deallocations for #iree_codegen.workgroup_local memory. The storage
+/// is owned by the HAL dispatch frame and released when the dispatch returns.
+struct ConvertWorkgroupLocalDeallocOp
+    : ConvertOpToLLVMWithABIPattern<memref::DeallocOp> {
+  ConvertWorkgroupLocalDeallocOp(HALDispatchABI &abi,
+                                 LLVMTypeConverter &typeConverter)
+      : ConvertOpToLLVMWithABIPattern(abi, typeConverter, /*benefit=*/2) {}
+
+  LogicalResult
+  matchAndRewrite(memref::DeallocOp deallocOp, memref::DeallocOpAdaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto memRefType = dyn_cast<MemRefType>(deallocOp.getMemref().getType());
+    if (!memRefType ||
+        !isa_and_nonnull<IREE::Codegen::WorkgroupLocalMemoryAttr>(
+            memRefType.getMemorySpace())) {
+      return failure();
+    }
+    rewriter.eraseOp(deallocOp);
+    return success();
+  }
+};
+
 struct InstrumentationEntry {
   // !llvm.ptr<i8> pointing at the base of the ringbuffer.
   Value basePtr;
@@ -1052,6 +1130,11 @@ void ConvertToLLVMPass::runOnOperation() {
   options.dataLayout = llvm::DataLayout(dataLayoutStr);
   options.overrideIndexBitwidth(options.dataLayout.getPointerSizeInBits());
   LLVMTypeConverter typeConverter(&getContext(), options, &dataLayoutAnalysis);
+  typeConverter.addTypeAttributeConversion(
+      [](BaseMemRefType, IREE::Codegen::WorkgroupLocalMemoryAttr attr)
+          -> TypeConverter::AttributeConversionResult {
+        return IntegerAttr::get(IntegerType::get(attr.getContext(), 64), 0);
+      });
 
   RewritePatternSet patterns(&getContext());
 
@@ -1123,6 +1206,8 @@ void ConvertToLLVMPass::runOnOperation() {
     ConvertHALInterfaceWorkgroupCountOp,
     ConvertHALInterfaceConstantLoadOp,
     ConvertHALInterfaceBindingSubspanOp,
+    ConvertWorkgroupLocalAllocOp,
+    ConvertWorkgroupLocalDeallocOp,
     ConvertHALInstrumentWorkgroupOp,
     ConvertHALInstrumentValueOp,
     ConvertHALInstrumentMemoryLoadOp,