[MLIR][Affine] Normalize memref.alloc ops with non trivial layout map #129875
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-memref Author: Arnab Dutta (arnab-polymage) ChangesPatch is 29.23 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/129875.diff 7 Files Affected:
diff --git a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
index 6ad39a3a91293..fefce0e0c087b 100644
--- a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
+++ b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
@@ -118,9 +118,11 @@ struct SemiAffineExprFlattener : public AffineExprFlattener {
// with a positive value." (see AffineExprKind in AffineExpr.h). If this
// assumption does not hold constraints (added above) are a contradiction.
+ return success();
+ } else if (localExpr.getKind() == AffineExprKind::Mul) {
+ (void)localVarCst.appendVar(VarKind::Local);
return success();
}
-
// TODO: Support other semi-affine expressions.
return failure();
}
@@ -163,7 +165,6 @@ getFlattenedAffineExprs(ArrayRef<AffineExpr> exprs, unsigned numDims,
return success();
};
-
if (addConservativeSemiAffineBounds) {
SemiAffineExprFlattener flattener(numDims, numSymbols);
return flattenExprs(flattener);
@@ -229,7 +230,8 @@ LogicalResult FlatLinearConstraints::composeMatchingMap(AffineMap other) {
assert(other.getNumSymbols() == getNumSymbolVars() && "symbol mismatch");
std::vector<SmallVector<int64_t, 8>> flatExprs;
- if (failed(flattenAlignedMapAndMergeLocals(other, &flatExprs)))
+ if (failed(flattenAlignedMapAndMergeLocals(
+ other, &flatExprs, /*addConservativeSemiAffineBounds=*/true)))
return failure();
assert(flatExprs.size() == other.getNumResults());
@@ -796,8 +798,6 @@ LogicalResult FlatLinearConstraints::flattenAlignedMapAndMergeLocals(
<< "composition unimplemented for semi-affine maps\n");
return failure();
}
-
- // Add localCst information.
if (localCst.getNumLocalVars() > 0) {
unsigned numLocalVars = getNumLocalVars();
// Insert local dims of localCst at the beginning.
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 7ef016f88be37..6671d981f2e4b 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1786,7 +1786,6 @@ static void createNewDynamicSizes(MemRefType oldMemRefType,
}
}
-// TODO: Currently works for static memrefs with a single layout map.
template <typename AllocLikeOp>
LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
MemRefType memrefType = allocOp->getType();
@@ -1799,7 +1798,6 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
// Either memrefType already had an identity map or the map couldn't be
// transformed to an identity map.
return failure();
-
Value oldMemRef = allocOp->getResult();
SmallVector<Value, 4> symbolOperands(allocOp->getSymbolOperands());
@@ -1819,8 +1817,40 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType, newDynamicSizes,
allocOp->getAlignmentAttr());
} else {
+ mlir::ValueRange dynamicSizes = allocOp->getDynamicSizes();
+ mlir::ValueRange symbolOperands = allocOp->getSymbolOperands();
+ ArrayRef<int64_t> newShape = newMemRefType.getShape();
+ ArrayRef<int64_t> oldShape = memrefType.getShape();
+ SmallVector<Value> mapOperands(oldShape.size() + symbolOperands.size());
+ SmallVector<Value> dimensionOperands;
+ unsigned dimId = 0, symId = 0;
+ // Collect all the map operands of `allocOp` (both dynamic sizes and symbol
+ // operands), which will help us to compute the dynamic sizes of the new
+ // alloc op we are going to create.
+ for (unsigned i = 0, e = oldShape.size(); i < e; i++) {
+ if (oldShape[i] == ShapedType::kDynamic)
+ mapOperands[i] = dynamicSizes[dimId++];
+ else
+ mapOperands[i] =
+ b.create<arith::ConstantIndexOp>(allocOp->getLoc(), oldShape[i]);
+ }
+ for (unsigned i = oldShape.size(), e = mapOperands.size(); i < e; i++)
+ mapOperands[i] = symbolOperands[symId++];
+ // Compute the dynamic sizes operands for the new alloc op. If `newShape` is
+ // dynamic along a dimension, compute its shape using the layout map and
+ // dynamic sizes and symbol operands of the old `allocOp`.
+ for (unsigned i = 0, e = newShape.size(); i < e; i++) {
+ if (newShape[i] != ShapedType::kDynamic)
+ continue;
+ AffineExpr resExpr = layoutMap.getResult(i);
+ auto resMap = AffineMap::get(layoutMap.getNumDims(),
+ layoutMap.getNumSymbols(), resExpr);
+ dimensionOperands.push_back(
+ b.create<AffineApplyOp>(allocOp->getLoc(), resMap, mapOperands));
+ }
newAlloc = b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType,
- allocOp->getAlignmentAttr());
+ dimensionOperands,
+ allocOp->getAlignmentAttr());
}
// Replace all uses of the old memref.
if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
@@ -1868,11 +1898,8 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
// Normalize only static memrefs and dynamic memrefs with a tiled-layout map
// for now.
- // TODO: Normalize the other types of dynamic memrefs.
SmallVector<std::tuple<AffineExpr, unsigned, unsigned>> tileSizePos;
(void)getTileSizePos(layoutMap, tileSizePos);
- if (memrefType.getNumDynamicDims() > 0 && tileSizePos.empty())
- return memrefType;
// We have a single map that is not an identity map. Create a new memref
// with the right shape and an identity layout map.
@@ -1894,7 +1921,6 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
unsigned newRank = layoutMap.getNumResults();
if (failed(fac.composeMatchingMap(layoutMap)))
return memrefType;
- // TODO: Handle semi-affine maps.
// Project out the old data dimensions.
fac.projectOut(newRank, fac.getNumVars() - newRank - fac.getNumLocalVars());
SmallVector<int64_t, 4> newShape(newRank);
@@ -1910,14 +1936,14 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
// For a static memref and an affine map with no symbols, this is
// always bounded. However, when we have symbols, we may not be able to
// obtain a constant upper bound. Also, mapping to a negative space is
- // invalid for normalization.
- if (!ubConst.has_value() || *ubConst < 0) {
- LLVM_DEBUG(llvm::dbgs()
- << "can't normalize map due to unknown/invalid upper bound");
- return memrefType;
- }
- // If dimension of new memrefType is dynamic, the value is -1.
- newShape[d] = *ubConst + 1;
+ // invalid for normalization. If dimension of new memrefType is dynamic,
+ // the value is `ShapedType::kDynamic`.
+ if (!ubConst.has_value())
+ newShape[d] = ShapedType::kDynamic;
+ else if (*ubConst >= 0)
+ newShape[d] = *ubConst + 1;
+ else
+ return memrefType;
}
// Create the new memref type after trivializing the old layout map.
diff --git a/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp b/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
index 08b853fe65b85..d2bd95b5996c8 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
@@ -445,8 +445,10 @@ void NormalizeMemRefs::normalizeFuncOpMemRefs(func::FuncOp funcOp,
if (oldMemRefType == newMemRefType)
continue;
// TODO: Assume single layout map. Multiple maps not supported.
+ // TODO: Semi-affine layout not supported.
AffineMap layoutMap = oldMemRefType.getLayout().getAffineMap();
- if (failed(replaceAllMemRefUsesWith(oldMemRef,
+ if (!layoutMap.getResult(0).isPureAffine() ||
+ failed(replaceAllMemRefUsesWith(oldMemRef,
/*newMemRef=*/newMemRef,
/*extraIndices=*/{},
/*indexRemap=*/layoutMap,
diff --git a/mlir/test/Dialect/Affine/memref-bound-check.mlir b/mlir/test/Dialect/Affine/memref-bound-check.mlir
index 80909abee51d6..321b2ba4a914f 100644
--- a/mlir/test/Dialect/Affine/memref-bound-check.mlir
+++ b/mlir/test/Dialect/Affine/memref-bound-check.mlir
@@ -124,13 +124,14 @@ func.func @mod_floordiv_nested() {
return
}
-// CHECK-LABEL: func @test_semi_affine_bailout
-func.func @test_semi_affine_bailout(%N : index) {
+// CHECK-LABEL: func @test_semi_affine_access
+func.func @test_semi_affine_access(%N : index) {
%B = memref.alloc() : memref<10 x i32>
affine.for %i = 0 to 10 {
%idx = affine.apply affine_map<(d0)[s0] -> (d0 * s0)>(%i)[%N]
%y = affine.load %B[%idx] : memref<10 x i32>
- // expected-error@-1 {{getMemRefRegion: compose affine map failed}}
+ // expected-error@-1 {{'affine.load' op memref out of upper bound access along dimension #1}}
+ // expected-error@-2 {{'affine.load' op memref out of lower bound access along dimension #1}}
}
return
}
diff --git a/mlir/test/Dialect/MemRef/normalize-memrefs-ops-dynamic.mlir b/mlir/test/Dialect/MemRef/normalize-memrefs-ops-dynamic.mlir
index a3f256b30c6a0..0cf8668561395 100644
--- a/mlir/test/Dialect/MemRef/normalize-memrefs-ops-dynamic.mlir
+++ b/mlir/test/Dialect/MemRef/normalize-memrefs-ops-dynamic.mlir
@@ -98,14 +98,15 @@ func.func @test_norm_dynamic1234(%arg0 : memref<?x?x?x?xf32, #map_tiled1>) -> ()
// -----
// Same test with maps that are not tiled layout maps in the arguments and the operations in the function.
-// This is not normalized since this is not tiled-layout map. No mod and floordiv.
#map_not_tiled0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2)>
-// CHECK-DAG: #[[$MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2)>
+// CHECK-DAG: #[[$MAP:.*]] = affine_map<(d0, d1, d2, d3) -> (d1)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d2 - d1)>
+// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3 - d2)>
// CHECK-LABEL: func @test_norm_dynamic_not_tiled0
-// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP6]]>) {
+// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x?xf32>) {
func.func @test_norm_dynamic_not_tiled0(%arg0 : memref<1x?x?x14xf32, #map_not_tiled0>) -> () {
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
@@ -118,11 +119,16 @@ func.func @test_norm_dynamic_not_tiled0(%arg0 : memref<1x?x?x14xf32, #map_not_ti
// CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index
// CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index
// CHECK-NOT: separator of consecutive DAGs
- // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
- // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
- // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP6]]>
- // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP6]]>, memref<1x?x?x14xf32, #[[$MAP6]]>) -> ()
- // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+ // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x?xf32>
+ // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x?xf32>
+ // CHECK-DAG: [[C_1_:%.+]] = arith.constant 1 : index
+ // CHECK-DAG: [[C_14_:%.+]] = arith.constant 14 : index
+ // CHECK: [[T0_:%.+]] = affine.apply #[[$MAP]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T1_:%.+]] = affine.apply #[[$MAP1]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T2_:%.+]] = affine.apply #[[$MAP2]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[RES_:%.+]] = memref.alloc([[T0_]], [[T1_]], [[T2_]]) : memref<1x?x?x?xf32>
+ // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x?xf32>, memref<1x?x?x?xf32>) -> ()
+ // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x?xf32>
// CHECK: return
}
@@ -133,10 +139,13 @@ func.func @test_norm_dynamic_not_tiled0(%arg0 : memref<1x?x?x14xf32, #map_not_ti
#map_not_tiled1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2, d2 mod 32, d3 mod 64)>
-// CHECK-DAG: #[[$MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2, d2 mod 32, d3 mod 64)>
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d2 - d1)>
+// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d3 - d2)>
+// CHECK-DAG: #[[$MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d2 mod 32)>
// CHECK-LABEL: func @test_norm_dynamic_not_tiled1
-// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP6]]>) {
+// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x?x?x64xf32>) {
func.func @test_norm_dynamic_not_tiled1(%arg0 : memref<1x?x?x14xf32, #map_not_tiled1>) -> () {
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
@@ -149,11 +158,17 @@ func.func @test_norm_dynamic_not_tiled1(%arg0 : memref<1x?x?x14xf32, #map_not_ti
// CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index
// CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index
// CHECK-NOT: separator of consecutive DAGs
- // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
- // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
- // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP6]]>
- // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP6]]>, memref<1x?x?x14xf32, #[[$MAP6]]>) -> ()
- // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+ // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x?x?x64xf32>
+ // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x?x?x64xf32>
+ // CHECK-DAG: [[C_1_:%.+]] = arith.constant 1 : index
+ // CHECK-DAG: [[C_14_:%.+]] = arith.constant 14 : index
+ // CHECK: [[T0_:%.+]] = affine.apply #[[$MAP]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T1_:%.+]] = affine.apply #[[$MAP1]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T2_:%.+]] = affine.apply #[[$MAP2]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T3_:%.+]] = affine.apply #[[$MAP3]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[RES_:%.+]] = memref.alloc([[T0_]], [[T1_]], [[T2_]], [[T3_]]) : memref<1x?x?x?x?x64xf32>
+ // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x?x?x64xf32>, memref<1x?x?x?x?x64xf32>) -> ()
+ // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x?x?x64xf32>
// CHECK: return
}
@@ -164,10 +179,12 @@ func.func @test_norm_dynamic_not_tiled1(%arg0 : memref<1x?x?x14xf32, #map_not_ti
#map_not_tiled2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 floordiv 64, d2 mod 32, d3 mod 32)>
-// CHECK-DAG: #[[$MAP7:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 floordiv 64, d2 mod 32, d3 mod 32)>
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d2 - d1)>
+// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d2 mod 32)>
// CHECK-LABEL: func @test_norm_dynamic_not_tiled2
-// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP7]]>) {
+// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x1x?x32xf32>) {
func.func @test_norm_dynamic_not_tiled2(%arg0 : memref<1x?x?x14xf32, #map_not_tiled2>) -> () {
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
@@ -180,11 +197,16 @@ func.func @test_norm_dynamic_not_tiled2(%arg0 : memref<1x?x?x14xf32, #map_not_ti
// CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index
// CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index
// CHECK-NOT: separator of consecutive DAGs
- // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP7]]>
- // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP7]]>
- // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP7]]>
- // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP7]]>, memref<1x?x?x14xf32, #[[$MAP7]]>) -> ()
- // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP7]]>
+ // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x1x?x32xf32>
+ // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x1x?x32xf32>
+ // CHECK-DAG: [[C_1_:%.+]] = arith.constant 1 : index
+ // CHECK-DAG: [[C_14_:%.+]] = arith.constant 14 : index
+ // CHECK: [[T0_:%.+]] = affine.apply #[[$MAP]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T1_:%.+]] = affine.apply #[[$MAP1]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T2_:%.+]] = affine.apply #[[$MAP2]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[RES_:%.+]] = memref.alloc([[T0_]], [[T1_]], [[T2_]]) : memref<1x?x?x1x?x32xf32>
+ // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x1x?x32xf32>, memref<1x?x?x1x?x32xf32>) -> ()
+ // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x1x?x32xf32>
// CHECK: return
}
@@ -195,10 +217,11 @@ func.func @test_norm_dynamic_not_tiled2(%arg0 : memref<1x?x?x14xf32, #map_not_ti
#map_not_tiled3 = affine_map<(d0, d1, d2, d3) -> (d0, d1 floordiv 32, d2, d3, d1 mod 32, d1 mod 32)>
-// CHECK-DAG: #[[$MAP8:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 floordiv 32, d2, d3, d1 mod 32, d1 mod 32)>
-
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1 floordiv 32)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d2)>
+// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d1 mod 32)>
// CHECK-LABEL: func @test_norm_dynamic_not_tiled3
-// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP8]]>) {
+// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14x?x?xf32>) {
func.func @test_norm_dynamic_not_tiled3(%arg0 : memref<1x?x?x14xf32, #map_not_tiled3>) -> () {
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
@@ -211,11 +234,17 @@ func.func @test_norm_dynamic_not_tiled3(%arg0 : memref<1x?x?x14xf32, #map_not_ti
// CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index
// CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index
// CHECK-NOT: separator of consecutive DAGs
- // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP8]]>
- // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP8]]>
- // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP8]]>
- // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP8]]>, memref<1x?x?x14xf32, #[[$MAP8]]>) -> ()
- // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP8]]>
+ // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14x?x?xf32>
+ // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14x?x?xf32>
+ // CHECK-DAG: [[C_1_:%.+]] = arith.constant 1 : index
+ // CHECK-DAG: [[C_14_:%.+]] = arith.constant 14 : index
+ // CHECK: [[T0_:%.+]] = affine.apply #[[$MAP]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T1_:%.+]] = affine.apply #[[$MAP1]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T2_:%.+]] = affine.apply #[[$MAP2]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T3_:%.+]]...
[truncated]
|
|
@llvm/pr-subscribers-mlir-affine Author: Arnab Dutta (arnab-polymage) ChangesPatch is 29.23 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/129875.diff 7 Files Affected:
diff --git a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
index 6ad39a3a91293..fefce0e0c087b 100644
--- a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
+++ b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
@@ -118,9 +118,11 @@ struct SemiAffineExprFlattener : public AffineExprFlattener {
// with a positive value." (see AffineExprKind in AffineExpr.h). If this
// assumption does not hold constraints (added above) are a contradiction.
+ return success();
+ } else if (localExpr.getKind() == AffineExprKind::Mul) {
+ (void)localVarCst.appendVar(VarKind::Local);
return success();
}
-
// TODO: Support other semi-affine expressions.
return failure();
}
@@ -163,7 +165,6 @@ getFlattenedAffineExprs(ArrayRef<AffineExpr> exprs, unsigned numDims,
return success();
};
-
if (addConservativeSemiAffineBounds) {
SemiAffineExprFlattener flattener(numDims, numSymbols);
return flattenExprs(flattener);
@@ -229,7 +230,8 @@ LogicalResult FlatLinearConstraints::composeMatchingMap(AffineMap other) {
assert(other.getNumSymbols() == getNumSymbolVars() && "symbol mismatch");
std::vector<SmallVector<int64_t, 8>> flatExprs;
- if (failed(flattenAlignedMapAndMergeLocals(other, &flatExprs)))
+ if (failed(flattenAlignedMapAndMergeLocals(
+ other, &flatExprs, /*addConservativeSemiAffineBounds=*/true)))
return failure();
assert(flatExprs.size() == other.getNumResults());
@@ -796,8 +798,6 @@ LogicalResult FlatLinearConstraints::flattenAlignedMapAndMergeLocals(
<< "composition unimplemented for semi-affine maps\n");
return failure();
}
-
- // Add localCst information.
if (localCst.getNumLocalVars() > 0) {
unsigned numLocalVars = getNumLocalVars();
// Insert local dims of localCst at the beginning.
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 7ef016f88be37..6671d981f2e4b 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1786,7 +1786,6 @@ static void createNewDynamicSizes(MemRefType oldMemRefType,
}
}
-// TODO: Currently works for static memrefs with a single layout map.
template <typename AllocLikeOp>
LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
MemRefType memrefType = allocOp->getType();
@@ -1799,7 +1798,6 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
// Either memrefType already had an identity map or the map couldn't be
// transformed to an identity map.
return failure();
-
Value oldMemRef = allocOp->getResult();
SmallVector<Value, 4> symbolOperands(allocOp->getSymbolOperands());
@@ -1819,8 +1817,40 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType, newDynamicSizes,
allocOp->getAlignmentAttr());
} else {
+ mlir::ValueRange dynamicSizes = allocOp->getDynamicSizes();
+ mlir::ValueRange symbolOperands = allocOp->getSymbolOperands();
+ ArrayRef<int64_t> newShape = newMemRefType.getShape();
+ ArrayRef<int64_t> oldShape = memrefType.getShape();
+ SmallVector<Value> mapOperands(oldShape.size() + symbolOperands.size());
+ SmallVector<Value> dimensionOperands;
+ unsigned dimId = 0, symId = 0;
+ // Collect all the map operands of `allocOp` (both dynamic sizes and symbol
+ // operands), which will help us to compute the dynamic sizes of the new
+ // alloc op we are going to create.
+ for (unsigned i = 0, e = oldShape.size(); i < e; i++) {
+ if (oldShape[i] == ShapedType::kDynamic)
+ mapOperands[i] = dynamicSizes[dimId++];
+ else
+ mapOperands[i] =
+ b.create<arith::ConstantIndexOp>(allocOp->getLoc(), oldShape[i]);
+ }
+ for (unsigned i = oldShape.size(), e = mapOperands.size(); i < e; i++)
+ mapOperands[i] = symbolOperands[symId++];
+ // Compute the dynamic sizes operands for the new alloc op. If `newShape` is
+ // dynamic along a dimension, compute its shape using the layout map and
+ // dynamic sizes and symbol operands of the old `allocOp`.
+ for (unsigned i = 0, e = newShape.size(); i < e; i++) {
+ if (newShape[i] != ShapedType::kDynamic)
+ continue;
+ AffineExpr resExpr = layoutMap.getResult(i);
+ auto resMap = AffineMap::get(layoutMap.getNumDims(),
+ layoutMap.getNumSymbols(), resExpr);
+ dimensionOperands.push_back(
+ b.create<AffineApplyOp>(allocOp->getLoc(), resMap, mapOperands));
+ }
newAlloc = b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType,
- allocOp->getAlignmentAttr());
+ dimensionOperands,
+ allocOp->getAlignmentAttr());
}
// Replace all uses of the old memref.
if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
@@ -1868,11 +1898,8 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
// Normalize only static memrefs and dynamic memrefs with a tiled-layout map
// for now.
- // TODO: Normalize the other types of dynamic memrefs.
SmallVector<std::tuple<AffineExpr, unsigned, unsigned>> tileSizePos;
(void)getTileSizePos(layoutMap, tileSizePos);
- if (memrefType.getNumDynamicDims() > 0 && tileSizePos.empty())
- return memrefType;
// We have a single map that is not an identity map. Create a new memref
// with the right shape and an identity layout map.
@@ -1894,7 +1921,6 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
unsigned newRank = layoutMap.getNumResults();
if (failed(fac.composeMatchingMap(layoutMap)))
return memrefType;
- // TODO: Handle semi-affine maps.
// Project out the old data dimensions.
fac.projectOut(newRank, fac.getNumVars() - newRank - fac.getNumLocalVars());
SmallVector<int64_t, 4> newShape(newRank);
@@ -1910,14 +1936,14 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
// For a static memref and an affine map with no symbols, this is
// always bounded. However, when we have symbols, we may not be able to
// obtain a constant upper bound. Also, mapping to a negative space is
- // invalid for normalization.
- if (!ubConst.has_value() || *ubConst < 0) {
- LLVM_DEBUG(llvm::dbgs()
- << "can't normalize map due to unknown/invalid upper bound");
- return memrefType;
- }
- // If dimension of new memrefType is dynamic, the value is -1.
- newShape[d] = *ubConst + 1;
+ // invalid for normalization. If dimension of new memrefType is dynamic,
+ // the value is `ShapedType::kDynamic`.
+ if (!ubConst.has_value())
+ newShape[d] = ShapedType::kDynamic;
+ else if (*ubConst >= 0)
+ newShape[d] = *ubConst + 1;
+ else
+ return memrefType;
}
// Create the new memref type after trivializing the old layout map.
diff --git a/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp b/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
index 08b853fe65b85..d2bd95b5996c8 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
@@ -445,8 +445,10 @@ void NormalizeMemRefs::normalizeFuncOpMemRefs(func::FuncOp funcOp,
if (oldMemRefType == newMemRefType)
continue;
// TODO: Assume single layout map. Multiple maps not supported.
+ // TODO: Semi-affine layout not supported.
AffineMap layoutMap = oldMemRefType.getLayout().getAffineMap();
- if (failed(replaceAllMemRefUsesWith(oldMemRef,
+ if (!layoutMap.getResult(0).isPureAffine() ||
+ failed(replaceAllMemRefUsesWith(oldMemRef,
/*newMemRef=*/newMemRef,
/*extraIndices=*/{},
/*indexRemap=*/layoutMap,
diff --git a/mlir/test/Dialect/Affine/memref-bound-check.mlir b/mlir/test/Dialect/Affine/memref-bound-check.mlir
index 80909abee51d6..321b2ba4a914f 100644
--- a/mlir/test/Dialect/Affine/memref-bound-check.mlir
+++ b/mlir/test/Dialect/Affine/memref-bound-check.mlir
@@ -124,13 +124,14 @@ func.func @mod_floordiv_nested() {
return
}
-// CHECK-LABEL: func @test_semi_affine_bailout
-func.func @test_semi_affine_bailout(%N : index) {
+// CHECK-LABEL: func @test_semi_affine_access
+func.func @test_semi_affine_access(%N : index) {
%B = memref.alloc() : memref<10 x i32>
affine.for %i = 0 to 10 {
%idx = affine.apply affine_map<(d0)[s0] -> (d0 * s0)>(%i)[%N]
%y = affine.load %B[%idx] : memref<10 x i32>
- // expected-error@-1 {{getMemRefRegion: compose affine map failed}}
+ // expected-error@-1 {{'affine.load' op memref out of upper bound access along dimension #1}}
+ // expected-error@-2 {{'affine.load' op memref out of lower bound access along dimension #1}}
}
return
}
diff --git a/mlir/test/Dialect/MemRef/normalize-memrefs-ops-dynamic.mlir b/mlir/test/Dialect/MemRef/normalize-memrefs-ops-dynamic.mlir
index a3f256b30c6a0..0cf8668561395 100644
--- a/mlir/test/Dialect/MemRef/normalize-memrefs-ops-dynamic.mlir
+++ b/mlir/test/Dialect/MemRef/normalize-memrefs-ops-dynamic.mlir
@@ -98,14 +98,15 @@ func.func @test_norm_dynamic1234(%arg0 : memref<?x?x?x?xf32, #map_tiled1>) -> ()
// -----
// Same test with maps that are not tiled layout maps in the arguments and the operations in the function.
-// This is not normalized since this is not tiled-layout map. No mod and floordiv.
#map_not_tiled0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2)>
-// CHECK-DAG: #[[$MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2)>
+// CHECK-DAG: #[[$MAP:.*]] = affine_map<(d0, d1, d2, d3) -> (d1)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d2 - d1)>
+// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3 - d2)>
// CHECK-LABEL: func @test_norm_dynamic_not_tiled0
-// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP6]]>) {
+// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x?xf32>) {
func.func @test_norm_dynamic_not_tiled0(%arg0 : memref<1x?x?x14xf32, #map_not_tiled0>) -> () {
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
@@ -118,11 +119,16 @@ func.func @test_norm_dynamic_not_tiled0(%arg0 : memref<1x?x?x14xf32, #map_not_ti
// CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index
// CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index
// CHECK-NOT: separator of consecutive DAGs
- // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
- // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
- // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP6]]>
- // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP6]]>, memref<1x?x?x14xf32, #[[$MAP6]]>) -> ()
- // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+ // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x?xf32>
+ // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x?xf32>
+ // CHECK-DAG: [[C_1_:%.+]] = arith.constant 1 : index
+ // CHECK-DAG: [[C_14_:%.+]] = arith.constant 14 : index
+ // CHECK: [[T0_:%.+]] = affine.apply #[[$MAP]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T1_:%.+]] = affine.apply #[[$MAP1]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T2_:%.+]] = affine.apply #[[$MAP2]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[RES_:%.+]] = memref.alloc([[T0_]], [[T1_]], [[T2_]]) : memref<1x?x?x?xf32>
+ // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x?xf32>, memref<1x?x?x?xf32>) -> ()
+ // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x?xf32>
// CHECK: return
}
@@ -133,10 +139,13 @@ func.func @test_norm_dynamic_not_tiled0(%arg0 : memref<1x?x?x14xf32, #map_not_ti
#map_not_tiled1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2, d2 mod 32, d3 mod 64)>
-// CHECK-DAG: #[[$MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 - d2, d2 mod 32, d3 mod 64)>
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d2 - d1)>
+// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d3 - d2)>
+// CHECK-DAG: #[[$MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d2 mod 32)>
// CHECK-LABEL: func @test_norm_dynamic_not_tiled1
-// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP6]]>) {
+// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x?x?x64xf32>) {
func.func @test_norm_dynamic_not_tiled1(%arg0 : memref<1x?x?x14xf32, #map_not_tiled1>) -> () {
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
@@ -149,11 +158,17 @@ func.func @test_norm_dynamic_not_tiled1(%arg0 : memref<1x?x?x14xf32, #map_not_ti
// CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index
// CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index
// CHECK-NOT: separator of consecutive DAGs
- // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
- // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
- // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP6]]>
- // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP6]]>, memref<1x?x?x14xf32, #[[$MAP6]]>) -> ()
- // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP6]]>
+ // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x?x?x64xf32>
+ // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x?x?x64xf32>
+ // CHECK-DAG: [[C_1_:%.+]] = arith.constant 1 : index
+ // CHECK-DAG: [[C_14_:%.+]] = arith.constant 14 : index
+ // CHECK: [[T0_:%.+]] = affine.apply #[[$MAP]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T1_:%.+]] = affine.apply #[[$MAP1]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T2_:%.+]] = affine.apply #[[$MAP2]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T3_:%.+]] = affine.apply #[[$MAP3]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[RES_:%.+]] = memref.alloc([[T0_]], [[T1_]], [[T2_]], [[T3_]]) : memref<1x?x?x?x?x64xf32>
+ // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x?x?x64xf32>, memref<1x?x?x?x?x64xf32>) -> ()
+ // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x?x?x64xf32>
// CHECK: return
}
@@ -164,10 +179,12 @@ func.func @test_norm_dynamic_not_tiled1(%arg0 : memref<1x?x?x14xf32, #map_not_ti
#map_not_tiled2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 floordiv 64, d2 mod 32, d3 mod 32)>
-// CHECK-DAG: #[[$MAP7:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 - d1, d3 floordiv 64, d2 mod 32, d3 mod 32)>
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d2 - d1)>
+// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d2 mod 32)>
// CHECK-LABEL: func @test_norm_dynamic_not_tiled2
-// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP7]]>) {
+// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x1x?x32xf32>) {
func.func @test_norm_dynamic_not_tiled2(%arg0 : memref<1x?x?x14xf32, #map_not_tiled2>) -> () {
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
@@ -180,11 +197,16 @@ func.func @test_norm_dynamic_not_tiled2(%arg0 : memref<1x?x?x14xf32, #map_not_ti
// CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index
// CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index
// CHECK-NOT: separator of consecutive DAGs
- // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP7]]>
- // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP7]]>
- // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP7]]>
- // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP7]]>, memref<1x?x?x14xf32, #[[$MAP7]]>) -> ()
- // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP7]]>
+ // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x1x?x32xf32>
+ // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x1x?x32xf32>
+ // CHECK-DAG: [[C_1_:%.+]] = arith.constant 1 : index
+ // CHECK-DAG: [[C_14_:%.+]] = arith.constant 14 : index
+ // CHECK: [[T0_:%.+]] = affine.apply #[[$MAP]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T1_:%.+]] = affine.apply #[[$MAP1]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T2_:%.+]] = affine.apply #[[$MAP2]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[RES_:%.+]] = memref.alloc([[T0_]], [[T1_]], [[T2_]]) : memref<1x?x?x1x?x32xf32>
+ // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x1x?x32xf32>, memref<1x?x?x1x?x32xf32>) -> ()
+ // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x1x?x32xf32>
// CHECK: return
}
@@ -195,10 +217,11 @@ func.func @test_norm_dynamic_not_tiled2(%arg0 : memref<1x?x?x14xf32, #map_not_ti
#map_not_tiled3 = affine_map<(d0, d1, d2, d3) -> (d0, d1 floordiv 32, d2, d3, d1 mod 32, d1 mod 32)>
-// CHECK-DAG: #[[$MAP8:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 floordiv 32, d2, d3, d1 mod 32, d1 mod 32)>
-
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1 floordiv 32)>
+// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d2)>
+// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d1 mod 32)>
// CHECK-LABEL: func @test_norm_dynamic_not_tiled3
-// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14xf32, #[[$MAP8]]>) {
+// CHECK-SAME: ([[ARG_0_:%.+]]: memref<1x?x?x14x?x?xf32>) {
func.func @test_norm_dynamic_not_tiled3(%arg0 : memref<1x?x?x14xf32, #map_not_tiled3>) -> () {
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
@@ -211,11 +234,17 @@ func.func @test_norm_dynamic_not_tiled3(%arg0 : memref<1x?x?x14xf32, #map_not_ti
// CHECK-DAG: [[CST_1_:%.+]] = arith.constant 1 : index
// CHECK-DAG: [[CST_2_:%.+]] = arith.constant 2 : index
// CHECK-NOT: separator of consecutive DAGs
- // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14xf32, #[[$MAP8]]>
- // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14xf32, #[[$MAP8]]>
- // CHECK: [[RES_:%.+]] = memref.alloc([[DIM_0_]], [[DIM_1_]]) : memref<1x?x?x14xf32, #[[$MAP8]]>
- // CHECK: "test.op_norm"([[ARG_0_]], [[RES_]]) : (memref<1x?x?x14xf32, #[[$MAP8]]>, memref<1x?x?x14xf32, #[[$MAP8]]>) -> ()
- // CHECK: memref.dealloc [[RES_]] : memref<1x?x?x14xf32, #[[$MAP8]]>
+ // CHECK-DAG: [[DIM_0_:%.+]] = memref.dim [[ARG_0_]], [[CST_1_]] : memref<1x?x?x14x?x?xf32>
+ // CHECK-DAG: [[DIM_1_:%.+]] = memref.dim [[ARG_0_]], [[CST_2_]] : memref<1x?x?x14x?x?xf32>
+ // CHECK-DAG: [[C_1_:%.+]] = arith.constant 1 : index
+ // CHECK-DAG: [[C_14_:%.+]] = arith.constant 14 : index
+ // CHECK: [[T0_:%.+]] = affine.apply #[[$MAP]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T1_:%.+]] = affine.apply #[[$MAP1]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T2_:%.+]] = affine.apply #[[$MAP2]]([[C_1_]], [[DIM_0_]], [[DIM_1_]], [[C_14_]])
+ // CHECK: [[T3_:%.+]]...
[truncated]
|
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
arnab-polymage
force-pushed
the
ornib/normalize_dynamic_alloc
branch
from
9f3b4a9 to
428e4e9
Compare
arnab-polymage
changed the title
bondhugula
reviewed
Apr 5, 2025
bondhugula
reviewed
Apr 5, 2025
| @@ -1819,8 +1817,40 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) { | |||
| b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType, newDynamicSizes, | |||
| allocOp->getAlignmentAttr()); | |||
| } else { | |||
| newAlloc = b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType, | |||
| allocOp->getAlignmentAttr()); | |||
| mlir::ValueRange dynamicSizes = allocOp->getDynamicSizes(); | |||
There was a problem hiding this comment.
Need a comment as to what this else block is meant to handle.
| // operands), which will help us to compute the dynamic sizes of the new | ||
| // alloc op we are going to create. | ||
| for (unsigned i = 0, e = oldShape.size(); i < e; i++) { | ||
| if (oldShape[i] == ShapedType::kDynamic) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Normalize memref.alloc ops with non trivial layout map