Skip to content

Commit 5c51165

Browse files
authored
[mlir][sparse] force a properly sized view on pos/crd/val under codegen (#91288)
Codegen "vectors" for pos/crd/val use the capacity as memref size, not the actual used size. Although the sparsifier itself always uses just the defined pos/crd/val parts, printing these and passing them back to a runtime environment could benefit from wrapping the basic pos/crd/val getters into a proper memref view that sets the right size.
1 parent 8bcb073 commit 5c51165

File tree

10 files changed

+240
-239
lines changed

10 files changed

+240
-239
lines changed

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

Lines changed: 30 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1050,10 +1050,14 @@ class SparseToPositionsConverter : public OpConversionPattern<ToPositionsOp> {
10501050
matchAndRewrite(ToPositionsOp op, OpAdaptor adaptor,
10511051
ConversionPatternRewriter &rewriter) const override {
10521052
// Replace the requested position access with corresponding field.
1053-
// The cast_op is inserted by type converter to intermix 1:N type
1054-
// conversion.
1053+
// The view is restricted to the actual size to ensure clients
1054+
// of this operation truly observe size, not capacity!
1055+
Location loc = op.getLoc();
1056+
Level lvl = op.getLevel();
10551057
auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
1056-
rewriter.replaceOp(op, desc.getPosMemRef(op.getLevel()));
1058+
auto mem = desc.getPosMemRef(lvl);
1059+
auto size = desc.getPosMemSize(rewriter, loc, lvl);
1060+
rewriter.replaceOp(op, genSliceToSize(rewriter, loc, mem, size));
10571061
return success();
10581062
}
10591063
};
@@ -1068,12 +1072,17 @@ class SparseToCoordinatesConverter
10681072
matchAndRewrite(ToCoordinatesOp op, OpAdaptor adaptor,
10691073
ConversionPatternRewriter &rewriter) const override {
10701074
// Replace the requested coordinates access with corresponding field.
1071-
// The cast_op is inserted by type converter to intermix 1:N type
1072-
// conversion.
1075+
// The view is restricted to the actual size to ensure clients
1076+
// of this operation truly observe size, not capacity!
1077+
Location loc = op.getLoc();
1078+
Level lvl = op.getLevel();
10731079
auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
1074-
rewriter.replaceOp(
1075-
op, desc.getCrdMemRefOrView(rewriter, op.getLoc(), op.getLevel()));
1076-
1080+
auto mem = desc.getCrdMemRefOrView(rewriter, loc, lvl);
1081+
if (lvl < getSparseTensorType(op.getTensor()).getAoSCOOStart()) {
1082+
auto size = desc.getCrdMemSize(rewriter, loc, lvl);
1083+
mem = genSliceToSize(rewriter, loc, mem, size);
1084+
}
1085+
rewriter.replaceOp(op, mem);
10771086
return success();
10781087
}
10791088
};
@@ -1088,11 +1097,14 @@ class SparseToCoordinatesBufferConverter
10881097
matchAndRewrite(ToCoordinatesBufferOp op, OpAdaptor adaptor,
10891098
ConversionPatternRewriter &rewriter) const override {
10901099
// Replace the requested coordinates access with corresponding field.
1091-
// The cast_op is inserted by type converter to intermix 1:N type
1092-
// conversion.
1100+
// The view is restricted to the actual size to ensure clients
1101+
// of this operation truly observe size, not capacity!
1102+
Location loc = op.getLoc();
1103+
Level lvl = getSparseTensorType(op.getTensor()).getAoSCOOStart();
10931104
auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
1094-
rewriter.replaceOp(op, desc.getAOSMemRef());
1095-
1105+
auto mem = desc.getAOSMemRef();
1106+
auto size = desc.getCrdMemSize(rewriter, loc, lvl);
1107+
rewriter.replaceOp(op, genSliceToSize(rewriter, loc, mem, size));
10961108
return success();
10971109
}
10981110
};
@@ -1106,10 +1118,13 @@ class SparseToValuesConverter : public OpConversionPattern<ToValuesOp> {
11061118
matchAndRewrite(ToValuesOp op, OpAdaptor adaptor,
11071119
ConversionPatternRewriter &rewriter) const override {
11081120
// Replace the requested values access with corresponding field.
1109-
// The cast_op is inserted by type converter to intermix 1:N type
1110-
// conversion.
1121+
// The view is restricted to the actual size to ensure clients
1122+
// of this operation truly observe size, not capacity!
1123+
Location loc = op.getLoc();
11111124
auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
1112-
rewriter.replaceOp(op, desc.getValMemRef());
1125+
auto mem = desc.getValMemRef();
1126+
auto size = desc.getValMemSize(rewriter, loc);
1127+
rewriter.replaceOp(op, genSliceToSize(rewriter, loc, mem, size));
11131128
return success();
11141129
}
11151130
};

mlir/test/Dialect/SparseTensor/binary_valued.mlir

Lines changed: 37 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,11 @@
2626
//
2727
// Make sure X += A * A => X += 1 in single loop.
2828
//
29-
//
3029
// CHECK-LABEL: func.func @sum_squares(
3130
// CHECK-SAME: %[[VAL_0:.*0]]: memref<?xindex>,
3231
// CHECK-SAME: %[[VAL_1:.*1]]: memref<?xindex>,
3332
// CHECK-SAME: %[[VAL_2:.*2]]: memref<?xf32>,
34-
// CHECK-SAME: %[[VAL_3:.*3]]: !sparse_tensor.storage_specifier<#{{.*}}>) -> memref<f32> {
33+
// CHECK-SAME: %[[VAL_3:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>) -> memref<f32> {
3534
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1.000000e+00 : f32
3635
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
3736
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index
@@ -40,23 +39,25 @@
4039
// CHECK-DAG: %[[VAL_9:.*]] = arith.constant 0.000000e+00 : f32
4140
// CHECK: %[[VAL_10:.*]] = memref.alloc() {alignment = 64 : i64} : memref<f32>
4241
// CHECK: linalg.fill ins(%[[VAL_9]] : f32) outs(%[[VAL_10]] : memref<f32>)
43-
// CHECK: %[[VAL_11:.*]] = memref.load %[[VAL_10]][] : memref<f32>
44-
// CHECK: %[[VAL_12:.*]] = scf.for %[[VAL_13:.*]] = %[[VAL_6]] to %[[VAL_8]] step %[[VAL_5]] iter_args(%[[VAL_14:.*]] = %[[VAL_11]]) -> (f32) {
45-
// CHECK: %[[VAL_15:.*]] = arith.muli %[[VAL_13]], %[[VAL_7]] : index
46-
// CHECK: %[[VAL_16:.*]] = scf.for %[[VAL_17:.*]] = %[[VAL_6]] to %[[VAL_7]] step %[[VAL_5]] iter_args(%[[VAL_18:.*]] = %[[VAL_14]]) -> (f32) {
47-
// CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_17]], %[[VAL_15]] : index
48-
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_0]]{{\[}}%[[VAL_19]]] : memref<?xindex>
49-
// CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_19]], %[[VAL_5]] : index
50-
// CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_0]]{{\[}}%[[VAL_21]]] : memref<?xindex>
51-
// CHECK: %[[VAL_23:.*]] = scf.for %[[VAL_24:.*]] = %[[VAL_20]] to %[[VAL_22]] step %[[VAL_5]] iter_args(%[[VAL_25:.*]] = %[[VAL_18]]) -> (f32) {
52-
// CHECK: %[[VAL_26:.*]] = arith.addf %[[VAL_25]], %[[VAL_4]] : f32
53-
// CHECK: scf.yield %[[VAL_26]] : f32
42+
// CHECK: %[[VAL_11:.*]] = sparse_tensor.storage_specifier.get %[[VAL_3]]
43+
// CHECK: %[[VAL_12:.*]] = memref.subview %[[VAL_0]][0] {{\[}}%[[VAL_11]]] [1] : memref<?xindex> to memref<?xindex>
44+
// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_10]][] : memref<f32>
45+
// CHECK: %[[VAL_14:.*]] = scf.for %[[VAL_15:.*]] = %[[VAL_6]] to %[[VAL_8]] step %[[VAL_5]] iter_args(%[[VAL_16:.*]] = %[[VAL_13]]) -> (f32) {
46+
// CHECK: %[[VAL_17:.*]] = arith.muli %[[VAL_15]], %[[VAL_7]] : index
47+
// CHECK: %[[VAL_18:.*]] = scf.for %[[VAL_19:.*]] = %[[VAL_6]] to %[[VAL_7]] step %[[VAL_5]] iter_args(%[[VAL_20:.*]] = %[[VAL_16]]) -> (f32) {
48+
// CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_19]], %[[VAL_17]] : index
49+
// CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_21]]] : memref<?xindex>
50+
// CHECK: %[[VAL_23:.*]] = arith.addi %[[VAL_21]], %[[VAL_5]] : index
51+
// CHECK: %[[VAL_24:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_23]]] : memref<?xindex>
52+
// CHECK: %[[VAL_25:.*]] = scf.for %[[VAL_26:.*]] = %[[VAL_22]] to %[[VAL_24]] step %[[VAL_5]] iter_args(%[[VAL_27:.*]] = %[[VAL_20]]) -> (f32) {
53+
// CHECK: %[[VAL_28:.*]] = arith.addf %[[VAL_27]], %[[VAL_4]] : f32
54+
// CHECK: scf.yield %[[VAL_28]] : f32
5455
// CHECK: } {"Emitted from" = "linalg.generic"}
55-
// CHECK: scf.yield %[[VAL_23]] : f32
56+
// CHECK: scf.yield %[[VAL_25]] : f32
5657
// CHECK: } {"Emitted from" = "linalg.generic"}
57-
// CHECK: scf.yield %[[VAL_16]] : f32
58+
// CHECK: scf.yield %[[VAL_18]] : f32
5859
// CHECK: } {"Emitted from" = "linalg.generic"}
59-
// CHECK: memref.store %[[VAL_12]], %[[VAL_10]][] : memref<f32>
60+
// CHECK: memref.store %[[VAL_14]], %[[VAL_10]][] : memref<f32>
6061
// CHECK: return %[[VAL_10]] : memref<f32>
6162
// CHECK: }
6263
//
@@ -99,25 +100,29 @@ func.func @sum_squares(%a: tensor<2x3x8xf32, #Sparse>) -> tensor<f32> {
99100
// CHECK-DAG: %[[VAL_9:.*]] = arith.constant 0.000000e+00 : f32
100101
// CHECK: %[[VAL_10:.*]] = memref.alloc() {alignment = 64 : i64} : memref<f32>
101102
// CHECK: linalg.fill ins(%[[VAL_9]] : f32) outs(%[[VAL_10]] : memref<f32>)
102-
// CHECK: %[[VAL_11:.*]] = memref.load %[[VAL_10]][] : memref<f32>
103-
// CHECK: %[[VAL_12:.*]] = scf.for %[[VAL_13:.*]] = %[[VAL_6]] to %[[VAL_8]] step %[[VAL_5]] iter_args(%[[VAL_14:.*]] = %[[VAL_11]]) -> (f32) {
104-
// CHECK: %[[VAL_15:.*]] = arith.muli %[[VAL_13]], %[[VAL_7]] : index
105-
// CHECK: %[[VAL_16:.*]] = scf.for %[[VAL_17:.*]] = %[[VAL_6]] to %[[VAL_7]] step %[[VAL_5]] iter_args(%[[VAL_18:.*]] = %[[VAL_14]]) -> (f32) {
106-
// CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_17]], %[[VAL_15]] : index
107-
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_0]]{{\[}}%[[VAL_19]]] : memref<?xindex>
108-
// CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_19]], %[[VAL_5]] : index
109-
// CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_0]]{{\[}}%[[VAL_21]]] : memref<?xindex>
110-
// CHECK: %[[VAL_23:.*]] = scf.for %[[VAL_24:.*]] = %[[VAL_20]] to %[[VAL_22]] step %[[VAL_5]] iter_args(%[[VAL_25:.*]] = %[[VAL_18]]) -> (f32) {
111-
// CHECK: %[[VAL_26:.*]] = memref.load %[[VAL_1]]{{\[}}%[[VAL_24]]] : memref<?xindex>
112-
// CHECK: %[[VAL_27:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_13]], %[[VAL_17]], %[[VAL_26]]] : memref<2x3x8xf32>
113-
// CHECK: %[[VAL_28:.*]] = arith.addf %[[VAL_27]], %[[VAL_25]] : f32
114-
// CHECK: scf.yield %[[VAL_28]] : f32
103+
// CHECK: %[[VAL_11:.*]] = sparse_tensor.storage_specifier.get %[[VAL_3]]
104+
// CHECK: %[[VAL_12:.*]] = memref.subview %[[VAL_0]][0] {{\[}}%[[VAL_11]]] [1] : memref<?xindex> to memref<?xindex>
105+
// CHECK: %[[VAL_13:.*]] = sparse_tensor.storage_specifier.get %[[VAL_3]]
106+
// CHECK: %[[VAL_14:.*]] = memref.subview %[[VAL_1]][0] {{\[}}%[[VAL_13]]] [1] : memref<?xindex> to memref<?xindex>
107+
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_10]][] : memref<f32>
108+
// CHECK: %[[VAL_16:.*]] = scf.for %[[VAL_17:.*]] = %[[VAL_6]] to %[[VAL_8]] step %[[VAL_5]] iter_args(%[[VAL_18:.*]] = %[[VAL_15]]) -> (f32) {
109+
// CHECK: %[[VAL_19:.*]] = arith.muli %[[VAL_17]], %[[VAL_7]] : index
110+
// CHECK: %[[VAL_20:.*]] = scf.for %[[VAL_21:.*]] = %[[VAL_6]] to %[[VAL_7]] step %[[VAL_5]] iter_args(%[[VAL_22:.*]] = %[[VAL_18]]) -> (f32) {
111+
// CHECK: %[[VAL_23:.*]] = arith.addi %[[VAL_21]], %[[VAL_19]] : index
112+
// CHECK: %[[VAL_24:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_23]]] : memref<?xindex>
113+
// CHECK: %[[VAL_25:.*]] = arith.addi %[[VAL_23]], %[[VAL_5]] : index
114+
// CHECK: %[[VAL_26:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_25]]] : memref<?xindex>
115+
// CHECK: %[[VAL_27:.*]] = scf.for %[[VAL_28:.*]] = %[[VAL_24]] to %[[VAL_26]] step %[[VAL_5]] iter_args(%[[VAL_29:.*]] = %[[VAL_22]]) -> (f32) {
116+
// CHECK: %[[VAL_30:.*]] = memref.load %[[VAL_14]]{{\[}}%[[VAL_28]]] : memref<?xindex>
117+
// CHECK: %[[VAL_31:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_17]], %[[VAL_21]], %[[VAL_30]]] : memref<2x3x8xf32>
118+
// CHECK: %[[VAL_32:.*]] = arith.addf %[[VAL_31]], %[[VAL_29]] : f32
119+
// CHECK: scf.yield %[[VAL_32]] : f32
115120
// CHECK: } {"Emitted from" = "linalg.generic"}
116-
// CHECK: scf.yield %[[VAL_23]] : f32
121+
// CHECK: scf.yield %[[VAL_27]] : f32
117122
// CHECK: } {"Emitted from" = "linalg.generic"}
118-
// CHECK: scf.yield %[[VAL_16]] : f32
123+
// CHECK: scf.yield %[[VAL_20]] : f32
119124
// CHECK: } {"Emitted from" = "linalg.generic"}
120-
// CHECK: memref.store %[[VAL_12]], %[[VAL_10]][] : memref<f32>
125+
// CHECK: memref.store %[[VAL_16]], %[[VAL_10]][] : memref<f32>
121126
// CHECK: return %[[VAL_10]] : memref<f32>
122127
// CHECK: }
123128
//

mlir/test/Dialect/SparseTensor/codegen.mlir

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -266,7 +266,9 @@ func.func @sparse_dense_3d_dyn(%arg0: tensor<?x?x?xf64, #Dense3D>) -> index {
266266
// CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
267267
// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
268268
// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier
269-
// CHECK: return %[[A2]] : memref<?xi32>
269+
// CHECK: %[[S:.*]] = sparse_tensor.storage_specifier.get %[[A5]] pos_mem_sz at 1
270+
// CHECK: %[[V:.*]] = memref.subview %[[A2]][0] [%[[S]]] [1]
271+
// CHECK: return %[[V]] : memref<?xi32>
270272
func.func @sparse_positions_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi32> {
271273
%0 = sparse_tensor.positions %arg0 { level = 1 : index } : tensor<?x?xf64, #DCSR> to memref<?xi32>
272274
return %0 : memref<?xi32>
@@ -279,7 +281,9 @@ func.func @sparse_positions_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi32>
279281
// CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
280282
// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
281283
// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier
282-
// CHECK: return %[[A3]] : memref<?xi64>
284+
// CHECK: %[[S:.*]] = sparse_tensor.storage_specifier.get %[[A5]] crd_mem_sz at 1
285+
// CHECK: %[[V:.*]] = memref.subview %[[A3]][0] [%[[S]]] [1]
286+
// CHECK: return %[[V]] : memref<?xi64>
283287
func.func @sparse_indices_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi64> {
284288
%0 = sparse_tensor.coordinates %arg0 { level = 1 : index } : tensor<?x?xf64, #DCSR> to memref<?xi64>
285289
return %0 : memref<?xi64>
@@ -292,7 +296,9 @@ func.func @sparse_indices_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi64> {
292296
// CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
293297
// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
294298
// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier
295-
// CHECK: return %[[A4]] : memref<?xf64>
299+
// CHECK: %[[S:.*]] = sparse_tensor.storage_specifier.get %[[A5]] val_mem_sz
300+
// CHECK: %[[V:.*]] = memref.subview %[[A4]][0] [%[[S]]] [1]
301+
// CHECK: return %[[V]] : memref<?xf64>
296302
func.func @sparse_values_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xf64> {
297303
%0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #DCSR> to memref<?xf64>
298304
return %0 : memref<?xf64>
@@ -305,13 +311,14 @@ func.func @sparse_values_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xf64> {
305311
// CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
306312
// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
307313
// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier
308-
// CHECK: return %[[A4]] : memref<?xf64>
314+
// CHECK: %[[S:.*]] = sparse_tensor.storage_specifier.get %[[A5]] val_mem_sz
315+
// CHECK: %[[V:.*]] = memref.subview %[[A4]][0] [%[[S]]] [1]
316+
// CHECK: return %[[V]] : memref<?xf64>
309317
func.func @sparse_values_coo(%arg0: tensor<?x?x?xf64, #ccoo>) -> memref<?xf64> {
310318
%0 = sparse_tensor.values %arg0 : tensor<?x?x?xf64, #ccoo> to memref<?xf64>
311319
return %0 : memref<?xf64>
312320
}
313321

314-
315322
// CHECK-LABEL: func.func @sparse_indices_coo(
316323
// CHECK-SAME: %[[A0:.*0]]: memref<?xindex>,
317324
// CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
@@ -320,7 +327,7 @@ func.func @sparse_values_coo(%arg0: tensor<?x?x?xf64, #ccoo>) -> memref<?xf64> {
320327
// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
321328
// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier
322329
// CHECK: %[[C2:.*]] = arith.constant 2 : index
323-
// CHECK: %[[S0:.*]] = sparse_tensor.storage_specifier.get %[[A5]] crd_mem_sz at 1
330+
// CHECK: %[[S0:.*]] = sparse_tensor.storage_specifier.get %[[A5]] crd_mem_sz at 1
324331
// CHECK: %[[S2:.*]] = arith.divui %[[S0]], %[[C2]] : index
325332
// CHECK: %[[R1:.*]] = memref.subview %[[A3]][0] {{\[}}%[[S2]]] [2] : memref<?xindex> to memref<?xindex, strided<[2]>>
326333
// CHECK: %[[R2:.*]] = memref.cast %[[R1]] : memref<?xindex, strided<[2]>> to memref<?xindex, strided<[?], offset: ?>>
@@ -337,7 +344,9 @@ func.func @sparse_indices_coo(%arg0: tensor<?x?x?xf64, #ccoo>) -> memref<?xindex
337344
// CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
338345
// CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
339346
// CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier
340-
// CHECK: return %[[A3]] : memref<?xindex>
347+
// CHECK: %[[S:.*]] = sparse_tensor.storage_specifier.get %[[A5]] crd_mem_sz at 1
348+
// CHECK: %[[V:.*]] = memref.subview %[[A3]][0] [%[[S]]] [1]
349+
// CHECK: return %[[V]] : memref<?xindex>
341350
func.func @sparse_indices_buffer_coo(%arg0: tensor<?x?x?xf64, #ccoo>) -> memref<?xindex> {
342351
%0 = sparse_tensor.coordinates_buffer %arg0 : tensor<?x?x?xf64, #ccoo> to memref<?xindex>
343352
return %0 : memref<?xindex>

0 commit comments

Comments
 (0)