已合并
[Vector] preserve strided layout when collapsing to 1D #2743
hujiajun创建于 20 天前
[Vector] preserve strided layout when collapsing to 1D #2743
已合并
共 2 个文件变更+50-15
| @@ -103,27 +103,35 @@ std::optional<Value> collapseAllDims(Value value, OpBuilder &builder, | |||
| 103 | return value; | 103 | return value; |
| 104 | } | 104 | } |
| 105 | 105 | ||
| 106 | - int64_t collapsedDimSize = 1; | 106 | + // Flattening requires static shapes; dynamic dims cannot be collapsed. |
| 107 | - bool isStatic = true; | 107 | + if (llvm::any_of(inputShape, |
| 108 | - for (int64_t dim : inputShape) { | 108 | + [](int64_t dim) { return ShapedType::isDynamic(dim); })) { |
| 109 | - if (dim == ShapedType::kDynamic) { | ||
| 110 | - collapsedDimSize = ShapedType::kDynamic; | ||
| 111 | - isStatic = false; | ||
| 112 | - break; | ||
| 113 | - } | ||
| 114 | - collapsedDimSize *= dim; | ||
| 115 | - } | ||
| 116 | - if (!isStatic) { | ||
| 117 | return std::nullopt; | 109 | return std::nullopt; |
| 118 | } | 110 | } |
| 119 | 111 | ||
| 120 | MLIRContext *ctx = builder.getContext(); | 112 | MLIRContext *ctx = builder.getContext(); |
| 121 | - auto resultType = | ||
| 122 | - mlir::MemRefType::get({collapsedDimSize}, inputType.getElementType()); | ||
| 123 | - Type typeWithAttr = setUBMemScopeAttr(resultType, ctx); | ||
| 124 | - | ||
| 125 | SmallVector<ReassociationIndices> reassociations = { | 113 | SmallVector<ReassociationIndices> reassociations = { |
| 126 | llvm::to_vector(llvm::seq<int64_t>(0, rank))}; | 114 | llvm::to_vector(llvm::seq<int64_t>(0, rank))}; |
| 115 | + // Reject non-contiguous strided sources: collapsing them is not guaranteed, | ||
| 116 | + // and computeCollapsedType asserts on the failure below. Identity layouts are | ||
| 117 | + // always collapsible, so this preserves the original behavior for the common | ||
| 118 | + // contiguous/identity cases. | ||
| 119 | + if (!memref::CollapseShapeOp::isGuaranteedCollapsible(inputType, | ||
| 120 | + reassociations)) { | ||
| 121 | + return std::nullopt; | ||
| 122 | + } | ||
| 123 | + // Preserve the strided layout and dynamic offset of the source (e.g. a | ||
| 124 | + // rank-reduced subview) when collapsing to 1D; a plain MemRefType::get() | ||
| 125 | + // would drop them and produce an invalid collapse_shape. | ||
| 126 | + auto resultType = | ||
| 127 | + memref::CollapseShapeOp::computeCollapsedType(inputType, reassociations); | ||
| 128 | + // computeCollapsedType already preserves the source address space. In the | ||
| 129 | + // normalize-vector pipeline the source is always UB (VF args are inferred UB | ||
| 130 | + // by InferHIVMMemScope, and OutlineAllocInVF guarantees no allocs in VFs), so | ||
| 131 | + // this call is effectively a no-op. Keep it as a defensive fallback for a | ||
| 132 | + // hypothetical source that has not yet had its memory scope inferred. | ||
| 133 | + Type typeWithAttr = setUBMemScopeAttr(resultType, ctx); | ||
| 134 | + | ||
| 127 | Value collapse = builder.create<memref::CollapseShapeOp>( | 135 | Value collapse = builder.create<memref::CollapseShapeOp>( |
| 128 | loc, typeWithAttr, value, reassociations); | 136 | loc, typeWithAttr, value, reassociations); |
| 129 | setUBMemScopeAttr(collapse, ctx); | 137 | setUBMemScopeAttr(collapse, ctx); |
| @@ -275,3 +275,30 @@ func.func @flatten_unit_subview_rank_reduce(%arg0: memref<1x4xi32, #hivm.address | |||
| 275 | vector.print %cast : vector<4xi32> | 275 | vector.print %cast : vector<4xi32> |
| 276 | return | 276 | return |
| 277 | } | 277 | } |
| 278 | + | ||
| 279 | +// ----- | ||
| 280 | +// CHECK-LABEL: func.func @merge_16x16_to_64x64_inverse_kernel_mix_aiv_outlined_merged_merged_vf_2( | ||
| 281 | +// CHECK-SAME: %[[ARG0:.*]]: memref<8x16xf32, #hivm.address_space<ub>> | ||
| 282 | +// CHECK-SAME: %[[ARG1:.*]]: memref<8x16xf32, strided<[16, 1], offset: ?>, #hivm.address_space<ub>> | ||
| 283 | +// CHECK: scf.for %[[ITER:.*]] = | ||
| 284 | +// CHECK-DAG: %[[SUBVIEW_R:.*]] = memref.subview %[[ARG0]][%[[ITER]], 0] {{.*}} memref<8x16xf32, #hivm.address_space<ub>> to memref<4x16xf32, strided<[16, 1], offset: ?>, #hivm.address_space<ub>> | ||
| 285 | +// CHECK-DAG: %[[READ:.*]] = vector.transfer_read %[[SUBVIEW_R]][%c0, %c0] | ||
| 286 | +// CHECK-DAG: %[[READ_CAST:.*]] = vector.shape_cast %[[READ]] : vector<4x16xf32> to vector<64xf32> | ||
| 287 | +// CHECK-DAG: %[[COLL:.*]] = memref.collapse_shape %[[ARG1]] {{.*}} memref<8x16xf32, strided<[16, 1], offset: ?>, #hivm.address_space<ub>> into memref<128xf32, strided<[1], offset: ?>, #hivm.address_space<ub>> | ||
| 288 | +// CHECK-DAG: %[[OFF:.*]] = arith.muli %[[ITER]], %c16 : index | ||
| 289 | +// CHECK-DAG: %[[SUBVIEW_W:.*]] = memref.subview %[[COLL]][%[[OFF]]] {{.*}} memref<128xf32, strided<[1], offset: ?>, #hivm.address_space<ub>> to memref<64xf32, strided<[1], offset: ?>, #hivm.address_space<ub>> | ||
| 290 | +// CHECK-DAG: vector.transfer_write %[[READ_CAST]], %[[SUBVIEW_W]][%c0] {{.*}} vector<64xf32>, memref<64xf32, strided<[1], offset: ?>, #hivm.address_space<ub>> | ||
| 291 | +// CHECK: return | ||
| 292 | +func.func @merge_16x16_to_64x64_inverse_kernel_mix_aiv_outlined_merged_merged_vf_2(%arg0: memref<8x16xf32, #hivm.address_space<ub>>, %arg1: memref<8x16xf32, strided<[16, 1], offset: ?>, #hivm.address_space<ub>>) attributes {hivm.func_core_type = #hivm.func_core_type<AIV>, hivm.storage_aligned, hivm.vector_function} { | ||
| 293 | + %c4 = arith.constant 4 : index | ||
| 294 | + %c8 = arith.constant 8 : index | ||
| 295 | + %c0 = arith.constant 0 : index | ||
| 296 | + %cst = arith.constant 0.000000e+00 : f32 | ||
| 297 | + scf.for %arg2 = %c0 to %c8 step %c4 { | ||
| 298 | + %subview_0 = memref.subview %arg0[%arg2, 0] [4, 16] [1, 1] : memref<8x16xf32, #hivm.address_space<ub>> to memref<4x16xf32, strided<[16, 1], offset: ?>, #hivm.address_space<ub>> | ||
| 299 | + %subview_2 = memref.subview %arg1[%arg2, 0] [4, 16] [1, 1] : memref<8x16xf32, strided<[16, 1], offset: ?>, #hivm.address_space<ub>> to memref<4x16xf32, strided<[16, 1], offset: ?>, #hivm.address_space<ub>> | ||
| 300 | + %0 = vector.transfer_read %subview_0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<4x16xf32, strided<[16, 1], offset: ?>, #hivm.address_space<ub>>, vector<4x16xf32> | ||
| 301 | + vector.transfer_write %0, %subview_2[%c0, %c0] {in_bounds = [true, true]} : vector<4x16xf32>, memref<4x16xf32, strided<[16, 1], offset: ?>, #hivm.address_space<ub>> | ||
| 302 | + } | ||
| 303 | + return | ||
| 304 | +} | ||