Skip to content

Commit c02aa91

Browse files
joker-ephjoker-eph
authored
Revert "[mlir][MemRef] Remove integer address space builders" (#138853)
Reverts #138579 An integration test is broken on the mlir-nvidia* bots.
1 parent 2f877c2 commit c02aa91

File tree

6 files changed

+100
-44
lines changed

6 files changed

+100
-44
lines changed

‎mlir/include/mlir/IR/BuiltinTypes.td

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -807,6 +807,11 @@ def Builtin_MemRef : Builtin_Type<"MemRef", "memref", [
807807
"ArrayRef<int64_t>":$shape, "Type":$elementType,
808808
CArg<"AffineMap">:$map,
809809
CArg<"Attribute", "{}">:$memorySpace)>,
810+
/// [deprecated] `Attribute`-based form should be used instead.
811+
TypeBuilderWithInferredContext<(ins
812+
"ArrayRef<int64_t>":$shape, "Type":$elementType,
813+
"AffineMap":$map,
814+
"unsigned":$memorySpaceInd)>
810815
];
811816
let extraClassDeclaration = [{
812817
using BaseMemRefType::clone;
@@ -1175,6 +1180,14 @@ def Builtin_UnrankedMemRef : Builtin_Type<"UnrankedMemRef", "unranked_memref", [
11751180
Attribute nonDefaultMemorySpace = skipDefaultMemorySpace(memorySpace);
11761181
return $_get(elementType.getContext(), elementType, nonDefaultMemorySpace);
11771182
}]>,
1183+
/// [deprecated] `Attribute`-based form should be used instead.
1184+
TypeBuilderWithInferredContext<(ins "Type":$elementType,
1185+
"unsigned":$memorySpace), [{
1186+
// Convert deprecated integer-like memory space to Attribute.
1187+
Attribute memorySpaceAttr =
1188+
wrapIntegerMemorySpace(memorySpace, elementType.getContext());
1189+
return UnrankedMemRefType::get(elementType, memorySpaceAttr);
1190+
}]>
11781191
];
11791192
let extraClassDeclaration = [{
11801193
using BaseMemRefType::clone;

‎mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2097,12 +2097,10 @@ static LogicalResult generateCopy(
20972097
// Check if a buffer was already created.
20982098
bool existingBuf = fastBufferMap.count(memref) > 0;
20992099
if (!existingBuf) {
2100-
Attribute fastMemorySpace;
2101-
if (copyOptions.fastMemorySpace != 0)
2102-
fastMemorySpace = prologue.getI64IntegerAttr(copyOptions.fastMemorySpace);
2100+
AffineMap fastBufferLayout = b.getMultiDimIdentityMap(rank);
21032101
auto fastMemRefType =
21042102
MemRefType::get(fastBufferShape, memRefType.getElementType(),
2105-
MemRefLayoutAttrInterface{}, fastMemorySpace);
2103+
fastBufferLayout, copyOptions.fastMemorySpace);
21062104

21072105
// Create the fast memory space buffer just before the 'affine.for'
21082106
// operation.
@@ -2177,12 +2175,8 @@ static LogicalResult generateCopy(
21772175
} else {
21782176
// DMA generation.
21792177
// Create a tag (single element 1-d memref) for the DMA.
2180-
Attribute tagMemorySpace;
2181-
if (copyOptions.tagMemorySpace != 0)
2182-
tagMemorySpace = prologue.getI64IntegerAttr(copyOptions.tagMemorySpace);
2183-
auto tagMemRefType =
2184-
MemRefType::get({1}, top.getIntegerType(32),
2185-
MemRefLayoutAttrInterface{}, tagMemorySpace);
2178+
auto tagMemRefType = MemRefType::get({1}, top.getIntegerType(32), {},
2179+
copyOptions.tagMemorySpace);
21862180
auto tagMemRef = prologue.create<memref::AllocOp>(loc, tagMemRefType);
21872181

21882182
SmallVector<Value, 4> tagIndices({zeroIndex});

‎mlir/lib/IR/BuiltinTypes.cpp

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -484,6 +484,14 @@ bool mlir::detail::isSupportedMemorySpace(Attribute memorySpace) {
484484
return false;
485485
}
486486

487+
Attribute mlir::detail::wrapIntegerMemorySpace(unsigned memorySpace,
488+
MLIRContext *ctx) {
489+
if (memorySpace == 0)
490+
return nullptr;
491+
492+
return IntegerAttr::get(IntegerType::get(ctx, 64), memorySpace);
493+
}
494+
487495
Attribute mlir::detail::skipDefaultMemorySpace(Attribute memorySpace) {
488496
IntegerAttr intMemorySpace = llvm::dyn_cast_or_null<IntegerAttr>(memorySpace);
489497
if (intMemorySpace && intMemorySpace.getValue() == 0)
@@ -575,6 +583,46 @@ MemRefType::getChecked(function_ref<InFlightDiagnostic()> emitErrorFn,
575583
elementType, layout, memorySpace);
576584
}
577585

586+
MemRefType MemRefType::get(ArrayRef<int64_t> shape, Type elementType,
587+
AffineMap map, unsigned memorySpaceInd) {
588+
589+
// Use default layout for empty map.
590+
if (!map)
591+
map = AffineMap::getMultiDimIdentityMap(shape.size(),
592+
elementType.getContext());
593+
594+
// Wrap AffineMap into Attribute.
595+
auto layout = AffineMapAttr::get(map);
596+
597+
// Convert deprecated integer-like memory space to Attribute.
598+
Attribute memorySpace =
599+
wrapIntegerMemorySpace(memorySpaceInd, elementType.getContext());
600+
601+
return Base::get(elementType.getContext(), shape, elementType, layout,
602+
memorySpace);
603+
}
604+
605+
MemRefType
606+
MemRefType::getChecked(function_ref<InFlightDiagnostic()> emitErrorFn,
607+
ArrayRef<int64_t> shape, Type elementType, AffineMap map,
608+
unsigned memorySpaceInd) {
609+
610+
// Use default layout for empty map.
611+
if (!map)
612+
map = AffineMap::getMultiDimIdentityMap(shape.size(),
613+
elementType.getContext());
614+
615+
// Wrap AffineMap into Attribute.
616+
auto layout = AffineMapAttr::get(map);
617+
618+
// Convert deprecated integer-like memory space to Attribute.
619+
Attribute memorySpace =
620+
wrapIntegerMemorySpace(memorySpaceInd, elementType.getContext());
621+
622+
return Base::getChecked(emitErrorFn, elementType.getContext(), shape,
623+
elementType, layout, memorySpace);
624+
}
625+
578626
LogicalResult MemRefType::verify(function_ref<InFlightDiagnostic()> emitError,
579627
ArrayRef<int64_t> shape, Type elementType,
580628
MemRefLayoutAttrInterface layout,

‎mlir/lib/IR/TypeDetail.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,9 @@ struct TupleTypeStorage final
140140
/// Checks if the memorySpace has supported Attribute type.
141141
bool isSupportedMemorySpace(Attribute memorySpace);
142142

143+
/// Wraps deprecated integer memory space to the new Attribute form.
144+
Attribute wrapIntegerMemorySpace(unsigned memorySpace, MLIRContext *ctx);
145+
143146
/// Replaces default memorySpace (integer == `0`) with empty Attribute.
144147
Attribute skipDefaultMemorySpace(Attribute memorySpace);
145148

‎mlir/test/Dialect/Vector/vector-warp-distribute.mlir

Lines changed: 28 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,10 @@
1919
// CHECK-SCF-IF-DAG: #[[$TIMES2:.*]] = affine_map<()[s0] -> (s0 * 2)>
2020
// CHECK-SCF-IF-DAG: #[[$TIMES4:.*]] = affine_map<()[s0] -> (s0 * 4)>
2121
// CHECK-SCF-IF-DAG: #[[$TIMES8:.*]] = affine_map<()[s0] -> (s0 * 8)>
22-
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_32xf32 : memref<32xf32, #gpu.address_space<workgroup>>
23-
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_64xf32 : memref<64xf32, #gpu.address_space<workgroup>>
24-
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_128xf32 : memref<128xf32, #gpu.address_space<workgroup>>
25-
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_256xf32 : memref<256xf32, #gpu.address_space<workgroup>>
22+
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_32xf32 : memref<32xf32, 3>
23+
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_64xf32 : memref<64xf32, 3>
24+
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_128xf32 : memref<128xf32, 3>
25+
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_256xf32 : memref<256xf32, 3>
2626

2727
// CHECK-SCF-IF-LABEL: func @rewrite_warp_op_to_scf_if(
2828
// CHECK-SCF-IF-SAME: %[[laneid:.*]]: index,
@@ -47,8 +47,8 @@ func.func @rewrite_warp_op_to_scf_if(%laneid: index,
4747
%r:2 = gpu.warp_execute_on_lane_0(%laneid)[32]
4848
args(%v0, %v1 : vector<4xf32>, vector<8xf32>) -> (vector<1xf32>, vector<2xf32>) {
4949
^bb0(%arg0: vector<128xf32>, %arg1: vector<256xf32>):
50-
// CHECK-SCF-IF: %[[arg1:.*]] = vector.transfer_read %[[buffer_v1]][%[[c0]]], %{{.*}} {in_bounds = [true]} : memref<256xf32, #gpu.address_space<workgroup>>, vector<256xf32>
51-
// CHECK-SCF-IF: %[[arg0:.*]] = vector.transfer_read %[[buffer_v0]][%[[c0]]], %{{.*}} {in_bounds = [true]} : memref<128xf32, #gpu.address_space<workgroup>>, vector<128xf32>
50+
// CHECK-SCF-IF: %[[arg1:.*]] = vector.transfer_read %[[buffer_v1]][%[[c0]]], %{{.*}} {in_bounds = [true]} : memref<256xf32, 3>, vector<256xf32>
51+
// CHECK-SCF-IF: %[[arg0:.*]] = vector.transfer_read %[[buffer_v0]][%[[c0]]], %{{.*}} {in_bounds = [true]} : memref<128xf32, 3>, vector<128xf32>
5252
// CHECK-SCF-IF: %[[def_0:.*]] = "some_def"(%[[arg0]]) : (vector<128xf32>) -> vector<32xf32>
5353
// CHECK-SCF-IF: %[[def_1:.*]] = "some_def"(%[[arg1]]) : (vector<256xf32>) -> vector<64xf32>
5454
%2 = "some_def"(%arg0) : (vector<128xf32>) -> vector<32xf32>
@@ -60,8 +60,8 @@ func.func @rewrite_warp_op_to_scf_if(%laneid: index,
6060
// CHECK-SCF-IF: }
6161
// CHECK-SCF-IF: gpu.barrier
6262
// CHECK-SCF-IF: %[[o1:.*]] = affine.apply #[[$TIMES2]]()[%[[laneid]]]
63-
// CHECK-SCF-IF: %[[r1:.*]] = vector.transfer_read %[[buffer_def_1]][%[[o1]]], %{{.*}} {in_bounds = [true]} : memref<64xf32, #gpu.address_space<workgroup>>, vector<2xf32>
64-
// CHECK-SCF-IF: %[[r0:.*]] = vector.transfer_read %[[buffer_def_0]][%[[laneid]]], %{{.*}} {in_bounds = [true]} : memref<32xf32, #gpu.address_space<workgroup>>, vector<1xf32>
63+
// CHECK-SCF-IF: %[[r1:.*]] = vector.transfer_read %[[buffer_def_1]][%[[o1]]], %{{.*}} {in_bounds = [true]} : memref<64xf32, 3>, vector<2xf32>
64+
// CHECK-SCF-IF: %[[r0:.*]] = vector.transfer_read %[[buffer_def_0]][%[[laneid]]], %{{.*}} {in_bounds = [true]} : memref<32xf32, 3>, vector<1xf32>
6565
// CHECK-SCF-IF: "some_use"(%[[r0]]) : (vector<1xf32>) -> ()
6666
// CHECK-SCF-IF: "some_use"(%[[r1]]) : (vector<2xf32>) -> ()
6767
"some_use"(%r#0) : (vector<1xf32>) -> ()
@@ -1065,18 +1065,18 @@ func.func @warp_execute_has_broadcast_semantics(%laneid: index, %s0: f32, %v0: v
10651065
args(%s0, %v0, %v1, %v2 : f32, vector<f32>, vector<1xf32>, vector<1x1xf32>) -> (f32, vector<f32>, vector<1xf32>, vector<1x1xf32>) {
10661066
^bb0(%bs0: f32, %bv0: vector<f32>, %bv1: vector<1xf32>, %bv2: vector<1x1xf32>):
10671067

1068-
// CHECK-SCF-IF: vector.transfer_read {{.*}}[%[[C0]], %[[C0]]]{{.*}} {in_bounds = [true, true]} : memref<1x1xf32, #gpu.address_space<workgroup>>, vector<1x1xf32>
1069-
// CHECK-SCF-IF: vector.transfer_read {{.*}}[%[[C0]]]{{.*}} {in_bounds = [true]} : memref<1xf32, #gpu.address_space<workgroup>>, vector<1xf32>
1070-
// CHECK-SCF-IF: vector.transfer_read {{.*}}[]{{.*}} : memref<f32, #gpu.address_space<workgroup>>, vector<f32>
1071-
// CHECK-SCF-IF: memref.load {{.*}}[%[[C0]]] : memref<1xf32, #gpu.address_space<workgroup>>
1068+
// CHECK-SCF-IF: vector.transfer_read {{.*}}[%[[C0]], %[[C0]]]{{.*}} {in_bounds = [true, true]} : memref<1x1xf32, 3>, vector<1x1xf32>
1069+
// CHECK-SCF-IF: vector.transfer_read {{.*}}[%[[C0]]]{{.*}} {in_bounds = [true]} : memref<1xf32, 3>, vector<1xf32>
1070+
// CHECK-SCF-IF: vector.transfer_read {{.*}}[]{{.*}} : memref<f32, 3>, vector<f32>
1071+
// CHECK-SCF-IF: memref.load {{.*}}[%[[C0]]] : memref<1xf32, 3>
10721072
// CHECK-SCF-IF: "some_def_0"(%{{.*}}) : (f32) -> f32
10731073
// CHECK-SCF-IF: "some_def_1"(%{{.*}}) : (vector<f32>) -> vector<f32>
10741074
// CHECK-SCF-IF: "some_def_1"(%{{.*}}) : (vector<1xf32>) -> vector<1xf32>
10751075
// CHECK-SCF-IF: "some_def_1"(%{{.*}}) : (vector<1x1xf32>) -> vector<1x1xf32>
1076-
// CHECK-SCF-IF: memref.store {{.*}}[%[[C0]]] : memref<1xf32, #gpu.address_space<workgroup>>
1077-
// CHECK-SCF-IF: vector.transfer_write {{.*}}[] : vector<f32>, memref<f32, #gpu.address_space<workgroup>>
1078-
// CHECK-SCF-IF: vector.transfer_write {{.*}}[%[[C0]]] {in_bounds = [true]} : vector<1xf32>, memref<1xf32, #gpu.address_space<workgroup>>
1079-
// CHECK-SCF-IF: vector.transfer_write {{.*}}[%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<1x1xf32>, memref<1x1xf32, #gpu.address_space<workgroup>>
1076+
// CHECK-SCF-IF: memref.store {{.*}}[%[[C0]]] : memref<1xf32, 3>
1077+
// CHECK-SCF-IF: vector.transfer_write {{.*}}[] : vector<f32>, memref<f32, 3>
1078+
// CHECK-SCF-IF: vector.transfer_write {{.*}}[%[[C0]]] {in_bounds = [true]} : vector<1xf32>, memref<1xf32, 3>
1079+
// CHECK-SCF-IF: vector.transfer_write {{.*}}[%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<1x1xf32>, memref<1x1xf32, 3>
10801080

10811081
%rs0 = "some_def_0"(%bs0) : (f32) -> f32
10821082
%rv0 = "some_def_1"(%bv0) : (vector<f32>) -> vector<f32>
@@ -1088,10 +1088,10 @@ func.func @warp_execute_has_broadcast_semantics(%laneid: index, %s0: f32, %v0: v
10881088
}
10891089

10901090
// CHECK-SCF-IF: gpu.barrier
1091-
// CHECK-SCF-IF: %[[RV2:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]]{{.*}} {in_bounds = [true, true]} : memref<1x1xf32, #gpu.address_space<workgroup>>, vector<1x1xf32>
1092-
// CHECK-SCF-IF: %[[RV1:.*]] = vector.transfer_read {{.*}}[%[[C0]]]{{.*}} {in_bounds = [true]} : memref<1xf32, #gpu.address_space<workgroup>>, vector<1xf32>
1093-
// CHECK-SCF-IF: %[[RV0:.*]] = vector.transfer_read {{.*}}[]{{.*}} : memref<f32, #gpu.address_space<workgroup>>, vector<f32>
1094-
// CHECK-SCF-IF: %[[RS0:.*]] = memref.load {{.*}}[%[[C0]]] : memref<1xf32, #gpu.address_space<workgroup>>
1091+
// CHECK-SCF-IF: %[[RV2:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]]{{.*}} {in_bounds = [true, true]} : memref<1x1xf32, 3>, vector<1x1xf32>
1092+
// CHECK-SCF-IF: %[[RV1:.*]] = vector.transfer_read {{.*}}[%[[C0]]]{{.*}} {in_bounds = [true]} : memref<1xf32, 3>, vector<1xf32>
1093+
// CHECK-SCF-IF: %[[RV0:.*]] = vector.transfer_read {{.*}}[]{{.*}} : memref<f32, 3>, vector<f32>
1094+
// CHECK-SCF-IF: %[[RS0:.*]] = memref.load {{.*}}[%[[C0]]] : memref<1xf32, 3>
10951095
// CHECK-SCF-IF: return %[[RS0]], %[[RV0]], %[[RV1]], %[[RV2]] : f32, vector<f32>, vector<1xf32>, vector<1x1xf32>
10961096
return %r#0, %r#1, %r#2, %r#3 : f32, vector<f32>, vector<1xf32>, vector<1x1xf32>
10971097
}
@@ -1106,22 +1106,22 @@ func.func @warp_execute_nd_distribute(%laneid: index, %v0: vector<1x64x1xf32>, %
11061106
-> (vector<1x64x1xf32>, vector<1x2x128xf32>) {
11071107
// CHECK-SCF-IF-DAG: %[[C0:.*]] = arith.constant 0 : index
11081108

1109-
// CHECK-SCF-IF: vector.transfer_write %{{.*}}, %{{.*}}[%[[LANEID]], %c0, %c0] {in_bounds = [true, true, true]} : vector<1x64x1xf32>, memref<32x64x1xf32, #gpu.address_space<workgroup>>
1109+
// CHECK-SCF-IF: vector.transfer_write %{{.*}}, %{{.*}}[%[[LANEID]], %c0, %c0] {in_bounds = [true, true, true]} : vector<1x64x1xf32>, memref<32x64x1xf32, 3>
11101110
// CHECK-SCF-IF: %[[RID:.*]] = affine.apply #[[$TIMES2]]()[%[[LANEID]]]
1111-
// CHECK-SCF-IF: vector.transfer_write %{{.*}}, %{{.*}}[%[[C0]], %[[RID]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x2x128xf32>, memref<1x64x128xf32, #gpu.address_space<workgroup>>
1111+
// CHECK-SCF-IF: vector.transfer_write %{{.*}}, %{{.*}}[%[[C0]], %[[RID]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x2x128xf32>, memref<1x64x128xf32, 3>
11121112
// CHECK-SCF-IF: gpu.barrier
11131113

11141114
// CHECK-SCF-IF: scf.if{{.*}}{
11151115
%r:2 = gpu.warp_execute_on_lane_0(%laneid)[32]
11161116
args(%v0, %v1 : vector<1x64x1xf32>, vector<1x2x128xf32>) -> (vector<1x64x1xf32>, vector<1x2x128xf32>) {
11171117
^bb0(%arg0: vector<32x64x1xf32>, %arg1: vector<1x64x128xf32>):
11181118

1119-
// CHECK-SCF-IF-DAG: %[[SR0:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : memref<32x64x1xf32, #gpu.address_space<workgroup>>, vector<32x64x1xf32>
1120-
// CHECK-SCF-IF-DAG: %[[SR1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : memref<1x64x128xf32, #gpu.address_space<workgroup>>, vector<1x64x128xf32>
1119+
// CHECK-SCF-IF-DAG: %[[SR0:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : memref<32x64x1xf32, 3>, vector<32x64x1xf32>
1120+
// CHECK-SCF-IF-DAG: %[[SR1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : memref<1x64x128xf32, 3>, vector<1x64x128xf32>
11211121
// CHECK-SCF-IF: %[[W0:.*]] = "some_def_0"(%[[SR0]]) : (vector<32x64x1xf32>) -> vector<32x64x1xf32>
11221122
// CHECK-SCF-IF: %[[W1:.*]] = "some_def_1"(%[[SR1]]) : (vector<1x64x128xf32>) -> vector<1x64x128xf32>
1123-
// CHECK-SCF-IF-DAG: vector.transfer_write %[[W0]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<32x64x1xf32>, memref<32x64x1xf32, #gpu.address_space<workgroup>>
1124-
// CHECK-SCF-IF-DAG: vector.transfer_write %[[W1]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x64x128xf32>, memref<1x64x128xf32, #gpu.address_space<workgroup>>
1123+
// CHECK-SCF-IF-DAG: vector.transfer_write %[[W0]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<32x64x1xf32>, memref<32x64x1xf32, 3>
1124+
// CHECK-SCF-IF-DAG: vector.transfer_write %[[W1]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x64x128xf32>, memref<1x64x128xf32, 3>
11251125

11261126
%r0 = "some_def_0"(%arg0) : (vector<32x64x1xf32>) -> vector<32x64x1xf32>
11271127
%r1 = "some_def_1"(%arg1) : (vector<1x64x128xf32>) -> vector<1x64x128xf32>
@@ -1132,8 +1132,8 @@ func.func @warp_execute_nd_distribute(%laneid: index, %v0: vector<1x64x1xf32>, %
11321132

11331133
// CHECK-SCF-IF: gpu.barrier
11341134
// CHECK-SCF-IF: %[[WID:.*]] = affine.apply #[[$TIMES2]]()[%[[LANEID]]]
1135-
// CHECK-SCF-IF-DAG: %[[R0:.*]] = vector.transfer_read %{{.*}}[%[[LANEID]], %[[C0]], %[[C0]]], %cst {in_bounds = [true, true, true]} : memref<32x64x1xf32, #gpu.address_space<workgroup>>, vector<1x64x1xf32>
1136-
// CHECK-SCF-IF-DAG: %[[R1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[WID]], %[[C0]]], %cst {in_bounds = [true, true, true]} : memref<1x64x128xf32, #gpu.address_space<workgroup>>, vector<1x2x128xf32>
1135+
// CHECK-SCF-IF-DAG: %[[R0:.*]] = vector.transfer_read %{{.*}}[%[[LANEID]], %[[C0]], %[[C0]]], %cst {in_bounds = [true, true, true]} : memref<32x64x1xf32, 3>, vector<1x64x1xf32>
1136+
// CHECK-SCF-IF-DAG: %[[R1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[WID]], %[[C0]]], %cst {in_bounds = [true, true, true]} : memref<1x64x128xf32, 3>, vector<1x2x128xf32>
11371137
// CHECK-SCF-IF: return %[[R0]], %[[R1]] : vector<1x64x1xf32>, vector<1x2x128xf32>
11381138
return %r#0, %r#1 : vector<1x64x1xf32>, vector<1x2x128xf32>
11391139
}

‎mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -518,17 +518,15 @@ struct TestVectorScanLowering
518518
static Value allocateGlobalSharedMemory(Location loc, OpBuilder &builder,
519519
gpu::WarpExecuteOnLane0Op warpOp,
520520
Type type) {
521-
Attribute sharedMemorySpaceAttr =
522-
builder.getAttr<gpu::AddressSpaceAttr>(gpu::AddressSpace::Workgroup);
521+
static constexpr int64_t kSharedMemorySpace = 3;
523522
// Compute type of shared memory buffer.
524523
MemRefType memrefType;
525524
if (auto vectorType = dyn_cast<VectorType>(type)) {
526525
memrefType =
527-
MemRefType::get(vectorType.getShape(), vectorType.getElementType(),
528-
MemRefLayoutAttrInterface{}, sharedMemorySpaceAttr);
526+
MemRefType::get(vectorType.getShape(), vectorType.getElementType(), {},
527+
kSharedMemorySpace);
529528
} else {
530-
memrefType = MemRefType::get({1}, type, MemRefLayoutAttrInterface{},
531-
sharedMemorySpaceAttr);
529+
memrefType = MemRefType::get({1}, type, {}, kSharedMemorySpace);
532530
}
533531

534532
// Get symbol table holding all shared memory globals.

0 commit comments

Comments
 (0)