folds memref.dim(gpu.alloc(%size), %idx) -> %size

[geexie]

Summary

implements simple canonicalization which eliminates memref.dim calls if possible

[Hardcode84]

I think, it can be added into GPUDialect.cpp instead of creating new file

[ftynse]

mlir/lib/Dialect/GPU/IR/GPUOps.cpp

1 ↗

(On Diff #369325)

FYI: this is not the full header, we also include the filename and a brief description in the top line.

17 ↗

(On Diff #369325)

Please do not commit commented-out code.

20–26 ↗

(On Diff #369325)

I'm pretty sure most of these are unnecessary. Please do not include files that are not used, this may increase compilation time significantly.

(Also, just put this code in GPUDialect.cpp as suggested above).

34 ↗

(On Diff #369325)

Nit: MLIR uses /// to document top-level entities such as classes and functions.

42 ↗

(On Diff #369325)

Nit: use named accessors instead of getOperand(1) with magic numbers, index() here.

52 ↗

(On Diff #369325)

Nit: please expand auto here. We use auto if the type is obvious from local context, e.g., there's a cast on the RHS, or annoyingly long to spell, e.g., iterators.

mlir/test/Dialect/GPU/canonicalize.mlir

25

Please add a newline

[Hardcode84]

mlir/lib/Dialect/GPU/IR/GPUOps.cpp

50 ↗

(On Diff #369325)

You can use getDefiningOp<AllocOp>() which will do dyn_cast for you

[geexie]

In D108892#2971965, @Hardcode84 wrote:

I think, it can be added into GPUDialect.cpp instead of creating new file

I saw in other dialects like memref that all foldings/canonicalization are placed into *Ops.cpp, while *Dialect contains basically only initialize method. I thought it's reasonable since it rains cpp files with td files. What do you think on this?

[ftynse]

In D108892#2972046, @geexie wrote:

In D108892#2971965, @Hardcode84 wrote:

I think, it can be added into GPUDialect.cpp instead of creating new file

I saw in other dialects like memref that all foldings/canonicalization are placed into *Ops.cpp, while *Dialect contains basically only initialize method. I thought it's reasonable since it rains cpp files with td files. What do you think on this?

We are rather inconsistent with naming, unfortunately. There are dialects that have NameOps.cpp, NameDialect.cpp or just Name.cpp that contain op implementations. Sometimes these files are also under the IR directory. Cleanups are always welcome as separate patches annotated as non-functional changes (NFC). For functional changes, we tend to prefer the minimally intrusive version. Since op implementations are already in GPUDialect.cpp, let's keep them there in this patch.

[geexie]

In D108892#2972048, @ftynse wrote:

In D108892#2972046, @geexie wrote:

In D108892#2971965, @Hardcode84 wrote:

I think, it can be added into GPUDialect.cpp instead of creating new file

I saw in other dialects like memref that all foldings/canonicalization are placed into *Ops.cpp, while *Dialect contains basically only initialize method. I thought it's reasonable since it rains cpp files with td files. What do you think on this?

We are rather inconsistent with naming, unfortunately. There are dialects that have NameOps.cpp, NameDialect.cpp or just Name.cpp that contain op implementations. Sometimes these files are also under the IR directory. Cleanups are always welcome as separate patches annotated as non-functional changes (NFC). For functional changes, we tend to prefer the minimally intrusive version. Since op implementations are already in GPUDialect.cpp, let's keep them there in this patch.

Got your point.

[geexie]

implements canonicalization which folds memref.dim(gpu.alloca(%size), %idx) -> %size

[geexie]

@ftynse, @Hardcode84, thanks for your suggestions. They are all addressed. Please proceed with the review

[ftynse]

Thanks!

[herhut]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

1118

Should this code check whether the size at index is actually dynamic?

This case is hard to trigger, as dim applied to a constant memref dimension is canonicalized by patterns for dim but still.

[Hardcode84]

nit: please fix alloca->alloc in description and comments

[geexie]

[mlir][gpu] folds memref.dim of gpu.alloc

implements canonicalization which folds memref.dim(gpu.alloc(%size), %idx) -> %size

[geexie]

In D108892#2972316, @Hardcode84 wrote:

nit: please fix alloca->alloc in description and comments

Done, thanks!

[geexie]

Thanks, @herhut for you suggestion. I added this check, hoverer appropriate test for this feature didn't come up to my mind. Do you have any ideas?

[herhut]

In D108892#2972360, @geexie wrote:

Thanks, @herhut for you suggestion. I added this check, hoverer appropriate test for this feature didn't come up to my mind. Do you have any ideas?

To test this, we would need to apply only the canonicalizations from the gpu dialect. The canonicalize pass does not support this, though. One could add support for limiting canonicalization to a specific dialect for testing purposes.

@ftynse WDYT?

I any case, I am happy with this to land as is.

[geexie]

[mlir][gpu] folds memref.dim of gpu.alloc

implements canonicalization which folds memref.dim(gpu.alloc(%size), %idx) -> %size

[geexie]

[mlir][gpu] folds memref.dim of gpu.alloc

implements canonicalization which folds memref.dim(gpu.alloc(%size), %idx) -> %size

[geexie]

In D108892#2972316, @Hardcode84 wrote:

nit: please fix alloca->alloc in description and comments

Finally done!

[ftynse]

In D108892#2972409, @herhut wrote:

In D108892#2972360, @geexie wrote:

Thanks, @herhut for you suggestion. I added this check, hoverer appropriate test for this feature didn't come up to my mind. Do you have any ideas?

To test this, we would need to apply only the canonicalizations from the gpu dialect. The canonicalize pass does not support this, though. One could add support for limiting canonicalization to a specific dialect for testing purposes.

@ftynse WDYT?

I any case, I am happy with this to land as is.

In the general scheme of things, I don't think even limiting to one dialect is sufficient, we may want to select specific patterns or, at least, a set of operations for which the patterns apply. There were also concerns about always having all patterns available in the pass that makes it expensive to run... Sounds like we should think about making the canonicalizer more flexible, but this is out of scope for this patch.

@geexie let me know if you need help landing this patch.

[rriddle]

In D108892#2974188, @ftynse wrote:

In D108892#2972409, @herhut wrote:

In D108892#2972360, @geexie wrote:

Thanks, @herhut for you suggestion. I added this check, hoverer appropriate test for this feature didn't come up to my mind. Do you have any ideas?

To test this, we would need to apply only the canonicalizations from the gpu dialect. The canonicalize pass does not support this, though. One could add support for limiting canonicalization to a specific dialect for testing purposes.

@ftynse WDYT?

I any case, I am happy with this to land as is.

In the general scheme of things, I don't think even limiting to one dialect is sufficient, we may want to select specific patterns or, at least, a set of operations for which the patterns apply. There were also concerns about always having all patterns available in the pass that makes it expensive to run... Sounds like we should think about making the canonicalizer more flexible, but this is out of scope for this patch.

The canonicalizer supports filtering patterns by label.

https://mlir.llvm.org/docs/PatternRewriter/#pattern-filtering-1
https://mlir.llvm.org/docs/PatternRewriter/#debug-names-and-labels

@geexie let me know if you need help landing this patch.

[geexie]

Thanks @ftynse! Just landed.

[geexie]

To test this, we would need to apply only the canonicalizations from the gpu dialect. The canonicalize pass does not support this, though. One could add support for limiting canonicalization to a specific dialect for testing purposes.

@herhut, that is a good idea. I might try implementing this.

In the general scheme of things, I don't think even limiting to one dialect is sufficient, we may want to select specific patterns or, at least, a set of operations for which the patterns apply. There were also concerns about always having all patterns available in the pass that makes it expensive to run... Sounds like we should think about making the canonicalizer more flexible, but this is out of scope for this patch.

@ftynse, I kind of understand that you mean. Do we need such functionality for no-testing use cases? do we have such level of customization in any of other passes?

The canonicalizer supports filtering patterns by label.

@rriddle, are those filters available in command line? We are kind of looking for a command

// RUN: mlir-opt %s -canonicalize,pattern=gpu --split-input-file -allow-unregistered-dialect | FileCheck %s

or something in this direction

[ftynse]

@ftynse, I kind of understand that you mean. Do we need such functionality for no-testing use cases? do we have such level of customization in any of other passes?

I don't think we have it elsewhere. The non-test use case for this is running the canonicalizer on IR where you know only a handful of dialects are used whereas much more are loaded in the context, and we therefore waste cycles on collecting and checking the patterns that we know cannot match. I don't have first-hand experience with this case so cannot tell how bad it actually is performance-wise.

[rriddle]

In D108892#2974304, @geexie wrote:

To test this, we would need to apply only the canonicalizations from the gpu dialect. The canonicalize pass does not support this, though. One could add support for limiting canonicalization to a specific dialect for testing purposes.

@herhut, that is a good idea. I might try implementing this.

In the general scheme of things, I don't think even limiting to one dialect is sufficient, we may want to select specific patterns or, at least, a set of operations for which the patterns apply. There were also concerns about always having all patterns available in the pass that makes it expensive to run... Sounds like we should think about making the canonicalizer more flexible, but this is out of scope for this patch.

@ftynse, I kind of understand that you mean. Do we need such functionality for no-testing use cases? do we have such level of customization in any of other passes?

The canonicalizer supports filtering patterns by label.

@rriddle, are those filters available in command line? We are kind of looking for a command

Yes, they are pass options on the canonicalizer (and any pattern pass that also hooks into those options).

https://github.com/llvm/llvm-project/blob/main/mlir/test/Transforms/test-canonicalize-filter.mlir

// RUN: mlir-opt %s -canonicalize,pattern=gpu --split-input-file -allow-unregistered-dialect | FileCheck %s

or something in this direction

[geexie]

In D108892#2975609, @rriddle wrote:

In D108892#2974304, @geexie wrote:

To test this, we would need to apply only the canonicalizations from the gpu dialect. The canonicalize pass does not support this, though. One could add support for limiting canonicalization to a specific dialect for testing purposes.

@herhut, that is a good idea. I might try implementing this.

In the general scheme of things, I don't think even limiting to one dialect is sufficient, we may want to select specific patterns or, at least, a set of operations for which the patterns apply. There were also concerns about always having all patterns available in the pass that makes it expensive to run... Sounds like we should think about making the canonicalizer more flexible, but this is out of scope for this patch.

@ftynse, I kind of understand that you mean. Do we need such functionality for no-testing use cases? do we have such level of customization in any of other passes?

The canonicalizer supports filtering patterns by label.

@rriddle, are those filters available in command line? We are kind of looking for a command

Yes, they are pass options on the canonicalizer (and any pattern pass that also hooks into those options).

https://github.com/llvm/llvm-project/blob/main/mlir/test/Transforms/test-canonicalize-filter.mlir

// RUN: mlir-opt %s -canonicalize,pattern=gpu --split-input-file -allow-unregistered-dialect | FileCheck %s

or something in this direction

@rriddle, thanks for a link. I didn’t know about this possibly and it looks the functionality I need. I’d try to write a test in similar way

[mehdi_amini]

To test this, we would need to apply only the canonicalizations from the gpu dialect.

Can you clarify what is "this" and why it requires to only apply a subset of the canonicalization patterns? This looks unusual to me.

[geexie]

In D108892#2975870, @mehdi_amini wrote:

To test this, we would need to apply only the canonicalizations from the gpu dialect.

Can you clarify what is "this" and why it requires to only apply a subset of the canonicalization patterns? This looks unusual to me.

@mehdi_amini, I added a check which tests if memrefType has a dynamic dimension in my pass since static dimensions got folded during memref::DimOp canonicalization. I do canonicalization for gpu::Alloc. With parametrized tests I can write a test which checks that memref.dim from constant dimension is not hit by -canonicalize pass for gpu dialect for which my pass is implemented.

mlir/include/mlir/Dialect/GPU/GPUOps.td

def GPU_AllocOp : GPU_Op<"alloc", [
   let extraClassDeclaration = [{
     MemRefType getType() { return memref().getType().cast<MemRefType>(); }
   }];
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) ` `
     `(` $dynamicSizes `)` (`` `[` $symbolOperands^ `]`)? attr-dict `:` type($memref)
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def GPU_DeallocOp : GPU_Op<"dealloc", [GPU_AsyncOpInterface]> {
 
   let summary = "GPU memory deallocation operation";
 
   let description = [{
     The `gpu.dealloc` operation frees the region of memory referenced by a
let description = [{
 
     The operation takes a scalar input and return a `!gpu.mma_matrix` where each
     element of is equal to the operand constant. The destination mma_matrix type
     must have elememt type equal to the constant type. Since the layout of
     `!gpu.mma_matrix` is opaque this only support setting all the elements to
     the same value.
 
     This op is meant to be used along with `gpu.subgroup_mma_compute`.
-    
     Example:
 
     ```mlir
      %0 = gpu.subgroup_mma_constant_matrix %a :
        !gpu.mma_matrix<16x16xf16, "AOp">
      %1 = gpu.subgroup_mma_constant_matrix %b :
        !gpu.mma_matrix<16x16xf32, "COp">
     ```

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

static LogicalResult foldMemRefCast(Operation *op) {
   return success(folded);
 }
 
 LogicalResult MemcpyOp::fold(ArrayRef<Attribute> operands,
                              SmallVectorImpl<::mlir::OpFoldResult> &results) {
   return foldMemRefCast(*this);
 }
 
+//===----------------------------------------------------------------------===//
+// GPU_AllocOp
+//===----------------------------------------------------------------------===//
+namespace {
+
+/// Folding of memref.dim(gpu.alloc(%size), %idx) -> %size similar to
+/// `memref::AllocOp`.
+struct SimplifyDimOfAllocOp : public OpRewritePattern<memref::DimOp> {
+  using OpRewritePattern<memref::DimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::DimOp dimOp,
+                                PatternRewriter &rewriter) const override {
+    auto index = dimOp.index().getDefiningOp<ConstantIndexOp>();
+    if (!index)
+      return failure();
+
+    auto memrefType = dimOp.source().getType().dyn_cast<MemRefType>();
+    if (!memrefType || !memrefType.isDynamicDim(index.getValue()))
+      return failure();
+
+    auto alloc = dimOp.source().getDefiningOp<AllocOp>();
+    if (!alloc)
+      return failure();
+
+    Value substituteOp = *(alloc.dynamicSizes().begin() +
+                           memrefType.getDynamicDimIndex(index.getValue()));
+    rewriter.replaceOp(dimOp, substituteOp);
+    return success();
+  }
+};
+
+} // end anonymous namespace.
+
+void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                          MLIRContext *context) {
+  results.add<SimplifyDimOfAllocOp>(context);
+}
+
 #include "mlir/Dialect/GPU/GPUOpInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"

mlir/test/Dialect/GPU/canonicalize.mlir

 // RUN: mlir-opt %s -canonicalize --split-input-file -allow-unregistered-dialect | FileCheck %s
 
 // CHECK-LABEL: @memcpy_after_cast
 func @memcpy_after_cast(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
   // CHECK-NOT: memref.cast
   // CHECK: gpu.memcpy
   %0 = memref.cast %arg0 : memref<10xf32> to memref<?xf32>
   %1 = memref.cast %arg1 : memref<10xf32> to memref<?xf32>
   gpu.memcpy %0,%1 : memref<?xf32>, memref<?xf32>
   return
 }
+
+// -----
+
+// Test case: Folding of memref.dim(gpu.alloc(%size), %idx) -> %size
+// CHECK-LABEL: func @gpu_dim_of_alloc(
+//  CHECK-SAME:     %[[SIZE:[0-9a-z]+]]: index
+//  CHECK-NEXT:   return %[[SIZE]] : index
+func @gpu_dim_of_alloc(%size: index) -> index {
+  %0 = gpu.alloc(%size) : memref<?xindex>
+  %c0 = constant 0 : index
+  %1 = memref.dim %0, %c0 : memref<?xindex>
+  return %1 : index
+}