[mlir][memref]: Fix Bug in GlobalOp Verifier

[FranklandJack]

When reconstructing the corresponding tensor type of a memref ensure we include the memory space of the tensor if it exists.

[llvmbot]

@llvm/pr-subscribers-mlir-memref

Author: Jack Frankland (FranklandJack)

Changes

When reconstructing the corresponding tensor type of a memref ensure we include the memory space of the tensor if it exists.

---

Full diff: https://github.com/llvm/llvm-project/pull/144900.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+2-1)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 11597505e7888..5db4ea30c25f7 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -59,7 +59,8 @@ LogicalResult mlir::memref::foldMemRefCast(Operation *op, Value inner) {
 /// type.
 Type mlir::memref::getTensorTypeFromMemRefType(Type type) {
   if (auto memref = llvm::dyn_cast<MemRefType>(type))
-    return RankedTensorType::get(memref.getShape(), memref.getElementType());
+    return RankedTensorType::get(memref.getShape(), memref.getElementType(),
+                                 memref.getMemorySpace());
   if (auto memref = llvm::dyn_cast<UnrankedMemRefType>(type))
     return UnrankedTensorType::get(memref.getElementType());
   return NoneType::get(type.getContext());

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Jack Frankland (FranklandJack)

Changes

When reconstructing the corresponding tensor type of a memref ensure we include the memory space of the tensor if it exists.

---

Full diff: https://github.com/llvm/llvm-project/pull/144900.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+2-1)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 11597505e7888..5db4ea30c25f7 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -59,7 +59,8 @@ LogicalResult mlir::memref::foldMemRefCast(Operation *op, Value inner) {
 /// type.
 Type mlir::memref::getTensorTypeFromMemRefType(Type type) {
   if (auto memref = llvm::dyn_cast<MemRefType>(type))
-    return RankedTensorType::get(memref.getShape(), memref.getElementType());
+    return RankedTensorType::get(memref.getShape(), memref.getElementType(),
+                                 memref.getMemorySpace());
   if (auto memref = llvm::dyn_cast<UnrankedMemRefType>(type))
     return UnrankedTensorType::get(memref.getElementType());
   return NoneType::get(type.getContext());

[banach-space]

Thanks! Please add a test :)

[FranklandJack]

Thanks! Please add a test :)

Yes, infact we can add two tests! One to check we can round trip a global memref with a tensor initializer that includes an encoding value and one to check that when the types mismatch we produce the correct error message.

Thanks for the review.

[banach-space]

OK, so this is setting the tensor encoding based on the memory space 🤔 This looks good and odd at the same time 😅

@matthias-springer , could you take a look with your expert eye? Based on -use-encoding-for-memory-space in https://mlir.llvm.org/docs/Passes/#-one-shot-bufferize, this would make sense. But it still feels like a very strong assumption, so not sure.

[FranklandJack]

Yes agreed this feels a bit odd so input on the best way forward very much appreciated here :D

Ultimately there is a question here about whether the encoding/address space of a tensor/memref needs to match its initializer. In the case of TOSA the answer is no (at least that isn’t something the tosa.const verifier checks), in the case of arith.const the answer seems to be yes, the operation inherits from an interface that requires they be the same.

For bufferization it seems like we should allow differences between the address space of the memref and the encoding of the tensor that initializes a global memref op since you may want to convert the memref’s address space like we do for spirv and we shouldn’t need to go back and update the initializer so that when the verifier does the type check they match.

[matthias-springer]

This PR assumes that tensor encoding and memref memory space are the same thing. That's not generally the case. Some compilers use that convention during bufferization, but not all do. E.g., the sparse compiler uses the tensor encoding for different purposes.

I don't quite follow what bug this PR is fixing. Why is it desirable that the tensor attribute encoding and the memref memory space match?

[FranklandJack]

This PR assumes that tensor encoding and memref memory space are the same thing. That's not generally the case. Some compilers use that convention during bufferization, but not all do. E.g., the sparse compiler uses the tensor encoding for different purposes.

I don't quite follow what bug this PR is fixing. Why is it desirable that the tensor attribute encoding and the memref memory space match?

Yeah sorry this issue here is a little bit subtle. To take a concerte example, we have an arith::ConstantOp which we want to attach an encoding to. Because arith::ConstantOp inherits the AllTypesMatch<["value", "result"] interface (or trait?) it needs the initializer type to match the value type of the returned tensor, so both need the encoding. During bufferization we pass the -use-encoding-for-memory-space and the arith::ConstantOp gets lowered to a memref::GlobalOp where the initializer is a tensor type with the encoding value and the operations results type is a memref with the corresponding memory space (these are both integers and happen to match when bufferizing with -use-encoding-for-memory-space).

The problem is that memref::GlobalOp::verify currently requires the types match but doesn't include the memory space :

// Check that the type of the initial value is compatible with the type of
    // the global variable.
    if (auto elementsAttr = llvm::dyn_cast<ElementsAttr>(initValue)) {
      Type initType = elementsAttr.getType();
      Type tensorType = getTensorTypeFromMemRefType(memrefType);
      if (initType != tensorType)
        return emitOpError("initial value expected to be of type ")
               << tensorType << ", but was of type " << initType;
    }
  }

So initType here includes an encoding and tensorType doesn't (even though the memrefType it mapped from does include an address space which happens to match the encoding in the case of -use-encoding-for-memory-space ).

As an example:

func.func @test_constant() -> () {
 %0 = "arith.constant"(){value = dense<[42, 24, 13, 52]> : tensor<4xi32, 1 : i32>} : () -> tensor<4xi32, 1 : i32>
 return
}

run with:

./build-debug/bin/mlir-opt --one-shot-bufferize="use-encoding-for-memory-space" example.mlir

with cause a verification error due to the situation described above.

I feel like there might be other solutions here:

1. Relax the memref::GlobalOp verifier to ignore the encoding when it compares the reconstructed tensor type against the initializer.
2. Relax the constraint that the types match exactly on arith::ContantOp and allow different encodings for the initializer and the returned value.

[matthias-springer]

I would change the memref.global verifier: it should check the shape + element type of the memref / tensor types match, nothing more. The only reason why there is a tensor type at all is because we want to reuse DenseElementsAttr, which requires a type.

[FranklandJack]

I would change the memref.global verifier: it should check the shape + element type of the memref / tensor types match, nothing more. The only reason why there is a tensor type at all is because we want to reuse DenseElementsAttr, which requires a type.

Sounds good to me. I've updated the verifier to do exactly that and added specific tests case for when the element types mismatch and when the shapes mismatch. I've also added a round trip to check we can now construct a memref in a given address space from an initializer with no matching encoding.

Let me know if you would like any further changes.

[banach-space]

Thanks, some suggestions inline.

[banach-space]

@FranklandJack , I noticed that you forced-pushed and erased the history. Please try to avoid that:

• https://llvm.org/docs/GitHub.html#rebasing-pull-requests-and-force-pushes

In particular, some earlier comments refer to the original revision, but that context is gone making it hard for folks referring to this in the future to follow the discussion.

Not a big deal, it is hard to keep track of all the LLVM policies/guidelines. And, TBH, many people do force-push (e.g. to rebase on top of main), but even then we should preserve the original commits.

Most importantly, thanks for fixing this!