[Mlir-commits] [mlir] [mlir] Make MemoryEffectsOpInterface optional. (PR #187427)

Slava Zakharin llvmlistbot at llvm.org
Thu Mar 19 12:54:46 PDT 2026


https://github.com/vzakhari updated https://github.com/llvm/llvm-project/pull/187427

>From b155ed70b0197135dfd9a6dbcb36db1d4d7371a9 Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Wed, 18 Mar 2026 18:51:10 -0700
Subject: [PATCH 1/2] [mlir] Make MemoryEffectsOpInterface optional.

For some instances of operations like calls it may be possible to report
specific memory effects and report "unknown" memory effects for other
instances. This patch makes MemoryEffectsOpInterface optional,
so that operations may implement it based on the particular instance.
---
 mlir/include/mlir-c/Interfaces.h              |  3 ++
 .../mlir/Interfaces/SideEffectInterfaces.td   | 23 +++++++++++++
 mlir/lib/CAPI/Interfaces/Interfaces.cpp       |  6 ++++
 .../Analysis/test-alias-analysis-modref.mlir  | 34 +++++++++++++++++++
 mlir/test/IR/test-side-effects.mlir           | 15 ++++++++
 mlir/test/lib/Dialect/Test/TestOpDefs.cpp     | 29 ++++++++++++++++
 mlir/test/lib/Dialect/Test/TestOps.td         |  7 ++++
 7 files changed, 117 insertions(+)

diff --git a/mlir/include/mlir-c/Interfaces.h b/mlir/include/mlir-c/Interfaces.h
index 17a812dcd86a9..5a4035f912319 100644
--- a/mlir/include/mlir-c/Interfaces.h
+++ b/mlir/include/mlir-c/Interfaces.h
@@ -113,6 +113,9 @@ typedef struct {
   /// Get memory effects callback.
   void (*getEffects)(MlirOperation op, MlirMemoryEffectInstancesList effects,
                      void *userData);
+  /// Callback that returns true iff the operation instance
+  /// can specify all its memory effects.
+  bool (*hasKnownMemoryEffects)(MlirOperation op);
   void *userData;
 } MlirMemoryEffectsOpInterfaceCallbacks;
 
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
index b292174fccb36..f8e1605556b97 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
@@ -31,6 +31,29 @@ def MemoryEffectsOpInterface
     an operation.
   }];
   let cppNamespace = "::mlir";
+
+  let append methods = [InterfaceMethod<
+                            [{
+      Returns true if the operation instance can specify all its memory effects.
+      The implementations may override this to make MemoryEffectsOpInterface
+      support conditional, e.g. when an operation may return conservative
+      "unknown" memory effects presenting itself as if it does not support
+      MemoryEffectsOpInterface at all.
+      By default, for all operations that are declared with
+      MemoryEffectsOpInterface support this method returns true.
+    }],
+                            /*retTy=*/"bool",
+                            /*methodName=*/"hasKnownMemoryEffects",
+                            /*args=*/(ins),
+                            /*methodBody=*/"",
+                            /*defaultImplementation=*/"return true;">,
+  ];
+
+  let extraClassOf = [{
+    // If hasKnownMemoryEffects() returns true, pretend that
+    // the operation instance does not implement MemoryEffectsOpInterface.
+    return $_op.hasKnownMemoryEffects();
+  }];
 }
 
 // The base class for defining specific memory effects.
diff --git a/mlir/lib/CAPI/Interfaces/Interfaces.cpp b/mlir/lib/CAPI/Interfaces/Interfaces.cpp
index 437a1dbab9dae..9a3cd6932cf5e 100644
--- a/mlir/lib/CAPI/Interfaces/Interfaces.cpp
+++ b/mlir/lib/CAPI/Interfaces/Interfaces.cpp
@@ -213,6 +213,12 @@ class MemoryEffectOpInterfaceFallbackModel
     callbacks.getEffects(wrap(op), cEffects, callbacks.userData);
   }
 
+  bool hasKnownMemoryEffects(Operation *op) const {
+    if (callbacks.hasKnownMemoryEffects)
+      return callbacks.hasKnownMemoryEffects(wrap(op));
+    return true;
+  }
+
 private:
   MlirMemoryEffectsOpInterfaceCallbacks callbacks;
 };
diff --git a/mlir/test/Analysis/test-alias-analysis-modref.mlir b/mlir/test/Analysis/test-alias-analysis-modref.mlir
index b82f789f4d77d..967974691dd58 100644
--- a/mlir/test/Analysis/test-alias-analysis-modref.mlir
+++ b/mlir/test/Analysis/test-alias-analysis-modref.mlir
@@ -89,3 +89,37 @@ func.func @addressable_write(%arg: memref<2xf32>) attributes {test.ptr = "func"}
   "test.side_effect_op"() {effects = [{effect="write"}], test.ptr = "side_effect_op"} : () -> i32
   return {test.ptr = "return"}
 }
+
+// -----
+
+// An op with known empty effects does not mod-ref any memory.
+// CHECK-LABEL: Testing : "optional_known_no_effects"
+// CHECK-DAG: optional_side_effect_op -> func.region0#0: NoModRef
+// CHECK-DAG: return -> func.region0#0: NoModRef
+func.func @optional_known_no_effects(%arg: memref<2xf32>) attributes {test.ptr = "func"} {
+  "test.optional_side_effect_op"() {effects = [], test.ptr = "optional_side_effect_op"} : () -> i32
+  return {test.ptr = "return"}
+}
+
+// -----
+
+// An op with unknown effects (no effects attr, hasKnownMemoryEffects = false)
+// is treated conservatively -> ModRef.
+// CHECK-LABEL: Testing : "optional_unknown_effects"
+// CHECK-DAG: optional_side_effect_op -> func.region0#0: ModRef
+// CHECK-DAG: return -> func.region0#0: NoModRef
+func.func @optional_unknown_effects(%arg: memref<2xf32>) attributes {test.ptr = "func"} {
+  "test.optional_side_effect_op"() {test.ptr = "optional_side_effect_op"} : () -> i32
+  return {test.ptr = "return"}
+}
+
+// -----
+
+// An op with a known read effect on an addressable resource -> Ref only.
+// CHECK-LABEL: Testing : "optional_known_read"
+// CHECK-DAG: optional_side_effect_op -> func.region0#0: Ref
+// CHECK-DAG: return -> func.region0#0: NoModRef
+func.func @optional_known_read(%arg: memref<2xf32>) attributes {test.ptr = "func"} {
+  "test.optional_side_effect_op"() {effects = [{effect="read"}], test.ptr = "optional_side_effect_op"} : () -> i32
+  return {test.ptr = "return"}
+}
diff --git a/mlir/test/IR/test-side-effects.mlir b/mlir/test/IR/test-side-effects.mlir
index b652ecb7dad1d..f04aeb9e6b696 100644
--- a/mlir/test/IR/test-side-effects.mlir
+++ b/mlir/test/IR/test-side-effects.mlir
@@ -56,5 +56,20 @@ func.func @side_effect(%arg : index) {
     test.region_yield %arg0 : index 
   } {effects = [ {effect="allocate", on_argument, test_resource} ]} : index -> index
 
+  // Known to have no memory effects (effects attr present but empty).
+  // expected-remark at +1 {{operation has no memory effects}}
+  %9 = "test.optional_side_effect_op"() {effects = []} : () -> i32
+
+  // Unknown effects (no effects attr) — the op does not cast to
+  // MemoryEffectOpInterface so the walk skips it; no diagnostic expected.
+  %10 = "test.optional_side_effect_op"() {} : () -> i32
+
+  // Known specific effects: read and write on <Default>.
+  // expected-remark at +2 {{found an instance of 'read' on resource '<Default>'}}
+  // expected-remark at +1 {{found an instance of 'write' on resource '<Default>'}}
+  %11 = "test.optional_side_effect_op"() {effects = [
+    {effect="read"}, {effect="write"}
+  ]} : () -> i32
+
   func.return 
 }
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index ed5dc5bead78a..5a36379bca1e1 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -501,6 +501,35 @@ void SideEffectOp::getEffects(
   testSideEffectOpGetEffect(getOperation(), effects);
 }
 
+//===----------------------------------------------------------------------===//
+// OptionalSideEffectOp
+//===----------------------------------------------------------------------===//
+
+bool OptionalSideEffectOp::hasKnownMemoryEffects() {
+  return (*this)->hasAttr("effects");
+}
+
+void OptionalSideEffectOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  ArrayAttr effectsAttr = (*this)->getAttrOfType<ArrayAttr>("effects");
+  if (!effectsAttr)
+    return;
+
+  for (Attribute element : effectsAttr) {
+    DictionaryAttr effectElement = cast<DictionaryAttr>(element);
+
+    MemoryEffects::Effect *effect =
+        StringSwitch<MemoryEffects::Effect *>(
+            cast<StringAttr>(effectElement.get("effect")).getValue())
+            .Case("allocate", MemoryEffects::Allocate::get())
+            .Case("free", MemoryEffects::Free::get())
+            .Case("read", MemoryEffects::Read::get())
+            .Case("write", MemoryEffects::Write::get());
+
+    effects.emplace_back(effect, SideEffects::DefaultResource::get());
+  }
+}
+
 void SideEffectWithRegionOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   // Check for an effects attribute on the op instance.
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 4c9e6b3fe9e45..298d7dcaa23f4 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2472,6 +2472,13 @@ def SideEffectOp : TEST_Op<"side_effect_op",
   let results = (outs AnyType:$result);
 }
 
+def OptionalSideEffectOp
+    : TEST_Op<"optional_side_effect_op",
+              [DeclareOpInterfaceMethods<
+                  MemoryEffectsOpInterface, ["hasKnownMemoryEffects"]>]> {
+  let results = (outs AnyType:$result);
+}
+
 def SideEffectWithRegionOp : TEST_Op<"side_effect_with_region_op",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      DeclareOpInterfaceMethods<TestEffectOpInterface>]> {

>From 633871a313a4ae781f7f6e704c3cefa449f3cd5c Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Thu, 19 Mar 2026 12:36:19 -0700
Subject: [PATCH 2/2] Added testing for C API.

---
 mlir/include/mlir-c/Interfaces.h              |  9 ++-
 .../mlir/Bindings/Python/IRInterfaces.h       | 11 ++-
 mlir/lib/Bindings/Python/IRInterfaces.cpp     | 24 ++++++
 mlir/lib/CAPI/Interfaces/CMakeLists.txt       |  3 +-
 mlir/lib/CAPI/Interfaces/Interfaces.cpp       |  6 +-
 .../dialects/memory_effects_interface.py      | 77 +++++++++++++++++++
 6 files changed, 124 insertions(+), 6 deletions(-)
 create mode 100644 mlir/test/python/dialects/memory_effects_interface.py

diff --git a/mlir/include/mlir-c/Interfaces.h b/mlir/include/mlir-c/Interfaces.h
index 5a4035f912319..914c7c2fdcf5d 100644
--- a/mlir/include/mlir-c/Interfaces.h
+++ b/mlir/include/mlir-c/Interfaces.h
@@ -104,6 +104,11 @@ mlirInferShapedTypeOpInterfaceInferReturnTypes(
 /// Returns the interface TypeID of the MemoryEffectsOpInterface.
 MLIR_CAPI_EXPORTED MlirTypeID mlirMemoryEffectsOpInterfaceTypeID(void);
 
+/// Returns true if the operation implements MemoryEffectsOpInterface, taking
+/// into account instance-level conditions (e.g. hasKnownMemoryEffects).
+MLIR_CAPI_EXPORTED bool
+mlirOperationImplementsMemoryEffectsOpInterface(MlirOperation operation);
+
 /// Callbacks for implementing MemoryEffectsOpInterface from external code.
 typedef struct {
   /// Optional constructor for user data. Set to nullptr to disable it.
@@ -114,8 +119,8 @@ typedef struct {
   void (*getEffects)(MlirOperation op, MlirMemoryEffectInstancesList effects,
                      void *userData);
   /// Callback that returns true iff the operation instance
-  /// can specify all its memory effects.
-  bool (*hasKnownMemoryEffects)(MlirOperation op);
+  /// can specify all its memory effects. Set to nullptr to always return true.
+  bool (*hasKnownMemoryEffects)(MlirOperation op, void *userData);
   void *userData;
 } MlirMemoryEffectsOpInterfaceCallbacks;
 
diff --git a/mlir/include/mlir/Bindings/Python/IRInterfaces.h b/mlir/include/mlir/Bindings/Python/IRInterfaces.h
index fb30e030b6c32..caebb2a1afae4 100644
--- a/mlir/include/mlir/Bindings/Python/IRInterfaces.h
+++ b/mlir/include/mlir/Bindings/Python/IRInterfaces.h
@@ -57,8 +57,7 @@ class PyConcreteOpInterface {
     }
 
     if (operation != nullptr) {
-      if (!mlirOperationImplementsInterface(*operation,
-                                            ConcreteIface::getInterfaceID())) {
+      if (!ConcreteIface::implementsInterface(*operation)) {
         std::string msg = "the operation does not implement ";
         throw nanobind::value_error((msg + ConcreteIface::pyClassName).c_str());
       }
@@ -101,6 +100,14 @@ class PyConcreteOpInterface {
   /// Hook for derived classes to add class-specific bindings.
   static void bindDerived(ClassTy &cls) {}
 
+  /// Returns true if the operation implements this interface. Derived classes
+  /// may override this to account for instance-level interface availability
+  /// (e.g., when extraClassOf logic makes the interface conditional).
+  static bool implementsInterface(MlirOperation operation) {
+    return mlirOperationImplementsInterface(operation,
+                                            ConcreteIface::getInterfaceID());
+  }
+
   /// Returns `true` if this object was constructed from a subclass of OpView
   /// rather than from an operation instance.
   bool isStatic() { return operation == nullptr; }
diff --git a/mlir/lib/Bindings/Python/IRInterfaces.cpp b/mlir/lib/Bindings/Python/IRInterfaces.cpp
index e7865cfda9d6f..8bfa975390c73 100644
--- a/mlir/lib/Bindings/Python/IRInterfaces.cpp
+++ b/mlir/lib/Bindings/Python/IRInterfaces.cpp
@@ -348,6 +348,12 @@ class PyMemoryEffectsOpInterface
   constexpr static GetTypeIDFunctionTy getInterfaceID =
       &mlirMemoryEffectsOpInterfaceTypeID;
 
+  /// Override to use isa<MemoryEffectOpInterface> which respects the
+  /// instance-level extraClassOf / hasKnownMemoryEffects check.
+  static bool implementsInterface(MlirOperation operation) {
+    return mlirOperationImplementsMemoryEffectsOpInterface(operation);
+  }
+
   /// Attach a new MemoryEffectsOpInterface FallbackModel to the named
   /// operation. The FallbackModel acts as a trampoline for callbacks on the
   /// Python class.
@@ -379,6 +385,24 @@ class PyMemoryEffectsOpInterface
       pyGetEffects(opview, effectsWrapper);
     };
 
+    // Set hasKnownMemoryEffects callback if the Python class defines it.
+    if (nb::hasattr(target, "has_known_memory_effects")) {
+      callbacks.hasKnownMemoryEffects = [](MlirOperation op,
+                                           void *userData) -> bool {
+        nb::handle pyClass(static_cast<PyObject *>(userData));
+        auto pyHasKnown = nb::cast<nb::callable>(
+            nb::getattr(pyClass, "has_known_memory_effects"));
+
+        PyMlirContextRef context =
+            PyMlirContext::forContext(mlirOperationGetContext(op));
+        auto opview = PyOperation::forOperation(context, op)->createOpView();
+
+        return nb::cast<bool>(pyHasKnown(opview));
+      };
+    } else {
+      callbacks.hasKnownMemoryEffects = nullptr;
+    }
+
     mlirMemoryEffectsOpInterfaceAttachFallbackModel(
         ctx->get(), mlirStringRefCreate(opName.c_str(), opName.size()),
         callbacks);
diff --git a/mlir/lib/CAPI/Interfaces/CMakeLists.txt b/mlir/lib/CAPI/Interfaces/CMakeLists.txt
index 7aefb56d992cf..dbe3f33975980 100644
--- a/mlir/lib/CAPI/Interfaces/CMakeLists.txt
+++ b/mlir/lib/CAPI/Interfaces/CMakeLists.txt
@@ -2,4 +2,5 @@ add_mlir_upstream_c_api_library(MLIRCAPIInterfaces
   Interfaces.cpp
 
   LINK_LIBS PUBLIC
-  MLIRInferTypeOpInterface)
+  MLIRInferTypeOpInterface
+  MLIRSideEffectInterfaces)
diff --git a/mlir/lib/CAPI/Interfaces/Interfaces.cpp b/mlir/lib/CAPI/Interfaces/Interfaces.cpp
index 9a3cd6932cf5e..5347436b02a7e 100644
--- a/mlir/lib/CAPI/Interfaces/Interfaces.cpp
+++ b/mlir/lib/CAPI/Interfaces/Interfaces.cpp
@@ -176,6 +176,10 @@ MlirTypeID mlirMemoryEffectsOpInterfaceTypeID() {
   return wrap(MemoryEffectOpInterface::getInterfaceID());
 }
 
+bool mlirOperationImplementsMemoryEffectsOpInterface(MlirOperation operation) {
+  return isa<MemoryEffectOpInterface>(unwrap(operation));
+}
+
 /// Fallback model for the MemoryEffectsOpInterface that uses C API callbacks.
 class MemoryEffectOpInterfaceFallbackModel
     : public mlir::MemoryEffectOpInterface::FallbackModel<
@@ -215,7 +219,7 @@ class MemoryEffectOpInterfaceFallbackModel
 
   bool hasKnownMemoryEffects(Operation *op) const {
     if (callbacks.hasKnownMemoryEffects)
-      return callbacks.hasKnownMemoryEffects(wrap(op));
+      return callbacks.hasKnownMemoryEffects(wrap(op), callbacks.userData);
     return true;
   }
 
diff --git a/mlir/test/python/dialects/memory_effects_interface.py b/mlir/test/python/dialects/memory_effects_interface.py
new file mode 100644
index 0000000000000..a48632ffc528e
--- /dev/null
+++ b/mlir/test/python/dialects/memory_effects_interface.py
@@ -0,0 +1,77 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir import ir
+from mlir.dialects import ext
+
+
+class TestMemEffectsDialect(ext.Dialect, name="test_mem_effects"):
+    pass
+
+
+class OptionalEffectsOp(TestMemEffectsDialect.Operation, name="optional_effects"):
+    pass
+
+
+def run(f):
+    print("\nTEST:", f.__name__)
+    f()
+    return f
+
+
+class OptionalMemoryEffectsModel(ir.MemoryEffectsOpInterface):
+    """MemoryEffectsOpInterface that is only active when 'effects' attr is present."""
+
+    @staticmethod
+    def has_known_memory_effects(op):
+        return "effects" in op.attributes
+
+    @staticmethod
+    def get_effects(op, effects):
+        pass
+
+
+def implements_memory_effects(op):
+    try:
+        ir.MemoryEffectsOpInterface(op)
+        return True
+    except ValueError:
+        return False
+
+
+# CHECK-LABEL: TEST: test_known_no_effects
+ at run
+def test_known_no_effects():
+    """When 'effects' attribute is present, hasKnownMemoryEffects returns true
+    and the interface is active."""
+    with ir.Context() as ctx, ir.Location.unknown():
+        TestMemEffectsDialect.load(reload=True)
+        OptionalMemoryEffectsModel.attach(
+            OptionalEffectsOp.OPERATION_NAME, context=ctx
+        )
+
+        module = ir.Module.create()
+        with ir.InsertionPoint(module.body):
+            op = ir.Operation.create("test_mem_effects.optional_effects")
+            op.attributes["effects"] = ir.UnitAttr.get()
+
+        # CHECK: implements MemoryEffectsOpInterface: True
+        print(f"implements MemoryEffectsOpInterface: {implements_memory_effects(op)}")
+
+
+# CHECK-LABEL: TEST: test_unknown_effects
+ at run
+def test_unknown_effects():
+    """When 'effects' attribute is absent, hasKnownMemoryEffects returns false
+    and the interface is not active (the op pretends not to implement it)."""
+    with ir.Context() as ctx, ir.Location.unknown():
+        TestMemEffectsDialect.load(reload=True)
+        OptionalMemoryEffectsModel.attach(
+            OptionalEffectsOp.OPERATION_NAME, context=ctx
+        )
+
+        module = ir.Module.create()
+        with ir.InsertionPoint(module.body):
+            op = ir.Operation.create("test_mem_effects.optional_effects")
+
+        # CHECK: implements MemoryEffectsOpInterface: False
+        print(f"implements MemoryEffectsOpInterface: {implements_memory_effects(op)}")



More information about the Mlir-commits mailing list