Index: mlir/lib/Dialect/Vector/VectorOps.cpp =================================================================== --- mlir/lib/Dialect/Vector/VectorOps.cpp +++ mlir/lib/Dialect/Vector/VectorOps.cpp @@ -2630,24 +2630,30 @@ template static LogicalResult foldTransferInBoundsAttribute(TransferOp op) { AffineMap permutationMap = op.permutation_map(); - if (!permutationMap.isMinorIdentity()) - return failure(); bool changed = false; SmallVector newInBounds; newInBounds.reserve(op.getTransferRank()); - op.zipResultAndIndexing([&](int64_t resultIdx, int64_t indicesIdx) { + for (int64_t resultIdx = 0, eResult = op.getTransferRank(); + resultIdx < eResult; ++resultIdx) { // Already marked as in-bounds, nothing to see here. if (op.isDimInBounds(resultIdx)) { newInBounds.push_back(true); - return; + continue; } + // Broadcassted dimensions are always in bound. + if (permutationMap.getResult(resultIdx).isa()) { + newInBounds.push_back(true); + changed = true; + continue; + } + int64_t indicesIdx = permutationMap.getDimPosition(resultIdx); // Currently out-of-bounds, check whether we can statically determine it is // inBounds. auto inBounds = isInBounds(op, resultIdx, indicesIdx); newInBounds.push_back(inBounds); // We commit the pattern if it is "more inbounds". changed |= inBounds; - }); + } if (!changed) return failure(); // OpBuilder is only used as a helper to build an I64ArrayAttr. Index: mlir/test/Dialect/Vector/canonicalize.mlir =================================================================== --- mlir/test/Dialect/Vector/canonicalize.mlir +++ mlir/test/Dialect/Vector/canonicalize.mlir @@ -534,8 +534,11 @@ // ----- +#map0 = affine_map<(d0, d1, d2, d3) -> (d3, 0, d1)> +#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d1)> + // CHECK-LABEL: fold_vector_transfers -func @fold_vector_transfers(%A: memref) -> (vector<4x8xf32>, vector<4x9xf32>) { +func @fold_vector_transfers(%A: memref, %B: memref<4x3x5x8xf32>) -> (vector<4x8xf32>, vector<4x9xf32>) { %c0 = constant 0 : index %f0 = constant 0.0 : f32 @@ -555,6 +558,12 @@ // CHECK-NOT: in_bounds vector.transfer_write %2, %A[%c0, %c0] : vector<4x9xf32>, memref + // CHECK: vector.transfer_read %{{.*}} {in_bounds = [false, true, true] + %3 = vector.transfer_read %B[%c0, %c0, %c0, %c0], %f0 {permutation_map = #map0} : memref<4x3x5x8xf32>, vector<9x12x3xf32> + + // CHECK: vector.transfer_write %{{.*}} {in_bounds = [false, false, true] + vector.transfer_write %3, %B[%c0, %c0, %c0, %c0] {permutation_map = #map1} : vector<9x12x3xf32>, memref<4x3x5x8xf32> + // CHECK: return return %1, %2 : vector<4x8xf32>, vector<4x9xf32> }