[Mlir-commits] [mlir] 0080d2a - [mlir][gpu] folds memref.dim of gpu.alloc

[marina kolpakova a.k.a. geexie]

Author: marina kolpakova a.k.a. geexie
Date: 2021-08-31T12:33:10+03:00
New Revision: 0080d2aa5542dc59c79c93b031ab58e1dd7306de

URL: https://github.com/llvm/llvm-project/commit/0080d2aa5542dc59c79c93b031ab58e1dd7306de
DIFF: https://github.com/llvm/llvm-project/commit/0080d2aa5542dc59c79c93b031ab58e1dd7306de.diff

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

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

Differential Revision: https://reviews.llvm.org/D108892

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/GPU/GPUOps.td
    mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
    mlir/test/Dialect/GPU/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/GPU/GPUOps.td b/mlir/include/mlir/Dialect/GPU/GPUOps.td
index 5eb0fbae0605..49c3a1dca5bf 100644
--- a/mlir/include/mlir/Dialect/GPU/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/GPUOps.td
@@ -834,6 +834,8 @@ def GPU_AllocOp : GPU_Op<"alloc", [
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) ` `
     `(` $dynamicSizes `)` (`` `[` $symbolOperands^ `]`)? attr-dict `:` type($memref)
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def GPU_DeallocOp : GPU_Op<"dealloc", [GPU_AsyncOpInterface]> {
@@ -1040,7 +1042,7 @@ def GPU_SubgroupMmaConstantMatrixOp : GPU_Op<"subgroup_mma_constant_matrix",
     the same value.
 
     This op is meant to be used along with `gpu.subgroup_mma_compute`.
-    
+
     Example:
 
     ```mlir

diff  --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 7f2622f5c707..8ec90cd5491d 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -1091,6 +1091,44 @@ LogicalResult MemcpyOp::fold(ArrayRef<Attribute> operands,
   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

diff  --git a/mlir/test/Dialect/GPU/canonicalize.mlir b/mlir/test/Dialect/GPU/canonicalize.mlir
index c3cba4719242..448e08be951b 100644
--- a/mlir/test/Dialect/GPU/canonicalize.mlir
+++ b/mlir/test/Dialect/GPU/canonicalize.mlir
@@ -9,3 +9,16 @@ func @memcpy_after_cast(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
   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
+}