Add catalyst.memspace attribute

doc/releases/changelog-dev.md

@@ -629,6 +629,11 @@
 
 <h3>Internal changes ⚙️</h3>
 
+* Added a generic memory-space attribute in the Catalyst dialect that can be used in the
+  `memory_space` slot of `memref` types (or as a discardable attribute) to tag values with
+  a logical memory domain plus an optional LLVM-level address space within that domain.
+  [(#2750)](https://github.com/PennyLaneAI/catalyst/pull/2750)
+
 * The compiler pipeline definitions now have a single source of truth. Previously, pipeline and
   pass sequences were duplicated between the frontend (`frontend/catalyst/pipelines.py`) and the
   compiler (`mlir/lib/Driver/Pipelines.cpp`). Now, there is a unique definition that lives in

mlir/include/Catalyst/IR/CMakeLists.txt

@@ -1,3 +1,9 @@
 add_mlir_dialect(CatalystOps catalyst)
 add_mlir_doc(CatalystDialect CatalystDialect Catalyst/ -gen-dialect-doc)
 add_mlir_doc(CatalystOps CatalystOps Catalyst/ -gen-op-doc)
+add_mlir_doc(CatalystAttributes CatalystAttributes Catalyst/ -gen-attrdef-doc)
+
+set(LLVM_TARGET_DEFINITIONS CatalystAttributes.td)
+mlir_tablegen(CatalystAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=catalyst)
+mlir_tablegen(CatalystAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=catalyst)
+add_public_tablegen_target(MLIRCatalystAttributesIncGen)

mlir/include/Catalyst/IR/CatalystAttributes.td

@@ -0,0 +1,67 @@
+// Copyright 2026 Xanadu Quantum Technologies Inc.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+//     http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef CATALYST_ATTRIBUTES
+#define CATALYST_ATTRIBUTES
+
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/BuiltinAttributes.td"
+include "mlir/IR/CommonAttrConstraints.td"
+
+include "Catalyst/IR/CatalystDialect.td"
+
+//===----------------------------------------------------------------------===//
+// Catalyst memory space attribute.
+//===----------------------------------------------------------------------===//
+
+class Catalyst_Attr<string name, string attrMnemonic, list<Trait> traits = [],
+                    string baseCppClass = "::mlir::Attribute">
+    : AttrDef<Catalyst_Dialect, name, traits, baseCppClass> {
+  let mnemonic = attrMnemonic;
+}
+
+def MemSpaceAttr : Catalyst_Attr<"MemSpace", "memspace"> {
+  let summary = "Catalyst memory-space attribute: memory domain plus optional address space.";
+  let description = [{
+    Generic memory-space attribute for use in the `memory_space` slot of
+    `memref` types (or as a discardable attribute in general) denoting a memory domain with an 
+    optional address space within that domain.
+
+    Syntax:
+
+    ```mlir
+    #catalyst.memspace<"<domain>">
+    #catalyst.memspace<"<domain>", <addr_space> : i64>
+    ```
+
+    Examples:
+
+    ```mlir
+    memref<4x64xi8, #catalyst.memspace<"d0">>
+    memref<1xi32,   #catalyst.memspace<"d1", 1 : i64>>
+    ```
+  }];
+
+  let parameters = (ins
+    StringRefParameter<"Symbolic memory domain name">:$domain,
+    OptionalParameter<"::mlir::IntegerAttr", "Optional explicit LLVM address space">:$addrSpace
+  );
+
+  let assemblyFormat =
+      "`<` $domain (`,` $addrSpace^)? `>`";
+
+  let genVerifyDecl = 1;
+}
+
+#endif // CATALYST_ATTRIBUTES

mlir/include/Catalyst/IR/CatalystDialect.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
 
@@ -30,3 +31,10 @@
 
 #define GET_TYPEDEF_CLASSES
 #include "Catalyst/IR/CatalystOpsTypes.h.inc"
+
+//===----------------------------------------------------------------------===//
+// Catalyst attribute declarations.
+//===----------------------------------------------------------------------===//
+
+#define GET_ATTRDEF_CLASSES
+#include "Catalyst/IR/CatalystAttributes.h.inc"

mlir/include/Catalyst/IR/CatalystDialect.td

@@ -39,6 +39,9 @@ def Catalyst_Dialect : Dialect {
 
     /// Use the default type printing/parsing hooks, otherwise we would explicitly define them.
     let useDefaultTypePrinterParser = 1;
+
+    /// Use the default attribute printer/parser hooks, otherwise we would explicitly define them.
+    let useDefaultAttributePrinterParser = 1;
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Catalyst/IR/CMakeLists.txt

@@ -7,6 +7,7 @@ add_mlir_library(MLIRCatalyst
 
     DEPENDS
     MLIRCatalystOpsIncGen
+    MLIRCatalystAttributesIncGen
 
     LINK_LIBS PRIVATE
 )

mlir/lib/Catalyst/IR/CatalystDialect.cpp

@@ -28,6 +28,9 @@ using namespace catalyst;
 
 #include "Catalyst/IR/CatalystOpsDialect.cpp.inc"
 
+#define GET_ATTRDEF_CLASSES
+#include "Catalyst/IR/CatalystAttributes.cpp.inc"
+
 //===----------------------------------------------------------------------===//
 // Catalyst dialect.
 //===----------------------------------------------------------------------===//
@@ -39,6 +42,11 @@ void CatalystDialect::initialize()
 #include "Catalyst/IR/CatalystOpsTypes.cpp.inc"
         >();
 
+    addAttributes<
+#define GET_ATTRDEF_LIST
+#include "Catalyst/IR/CatalystAttributes.cpp.inc"
+        >();
+
     addOperations<
 #define GET_OP_LIST
 #include "Catalyst/IR/CatalystOps.cpp.inc"
@@ -48,6 +56,30 @@ void CatalystDialect::initialize()
                               CallbackCallOp, CallbackOp>();
 }
 
+//===----------------------------------------------------------------------===//
+// MemSpaceAttr
+//===----------------------------------------------------------------------===//
+
+LogicalResult MemSpaceAttr::verify(function_ref<InFlightDiagnostic()> emitError,
+                                   llvm::StringRef domain, IntegerAttr addrSpace)
+{
+    if (domain.empty()) {
+        return emitError() << "catalyst.memspace: `domain` must be non-empty";
+    }
+
+    if (addrSpace) {
+        if (!addrSpace.getType().isSignlessInteger()) {
+            return emitError()
+                   << "catalyst.memspace: addr_space must be a signless integer attribute";
+        }
+        if (addrSpace.getValue().isNegative()) {
+            return emitError() << "catalyst.memspace: addr_space must be non-negative";
+        }
+    }
+
+    return success();
+}
+
 //===----------------------------------------------------------------------===//
 // CallbackOp
 //===----------------------------------------------------------------------===//

mlir/test/Catalyst/MemSpaceAttributeTest.mlir

@@ -0,0 +1,53 @@
+// Copyright 2026 Xanadu Quantum Technologies Inc.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+//     http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// RUN: quantum-opt %s --split-input-file --verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL: func.func @memspace_in_memref
+// CHECK-SAME:    (%arg0: memref<4x64xi8, #catalyst.memspace<"d0">>
+// CHECK-SAME:     %arg1: memref<1xi32, #catalyst.memspace<"d1", 1 : i64>>)
+func.func @memspace_in_memref(
+    %arg0: memref<4x64xi8, #catalyst.memspace<"d0">>,
+    %arg1: memref<1xi32,   #catalyst.memspace<"d1", 1 : i64>>
+) {
+    return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @memspace_on_op()
+// CHECK-SAME:    attributes {catalyst.memspace = #catalyst.memspace<"d0", 3 : i64>}
+func.func @memspace_on_op() attributes {
+    catalyst.memspace = #catalyst.memspace<"d0", 3 : i64>
+} {
+    return
+}
+
+// -----
+
+func.func @bad_domain(
+    // expected-error@+1 {{catalyst.memspace: `domain` must be non-empty}}
+    %arg0: memref<4xi8, #catalyst.memspace<"">>
+) {
+    return
+}
+
+// -----
+
+func.func @bad_addr_space(
+    // expected-error@+1 {{catalyst.memspace: addr_space must be non-negative}}
+    %arg0: memref<4xi8, #catalyst.memspace<"d0", -1 : i64>>
+) {
+    return
+}