diff --git a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp --- a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp +++ b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp @@ -1247,8 +1247,13 @@ Location loc = xferOp.getLoc(); auto vecType = xferOp.getVectorType(); auto lb = rewriter.create(loc, 0); - auto ub = + Value ub = rewriter.create(loc, vecType.getDimSize(0)); + if (vecType.isScalable()) { + Value vscale = + rewriter.create(loc, rewriter.getIndexType()); + ub = rewriter.create(loc, ub, vscale); + } auto step = rewriter.create(loc, 1); auto loopState = Strategy1d::initialLoopState(rewriter, xferOp); diff --git a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir --- a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir +++ b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir @@ -511,3 +511,39 @@ tensor, vector<1xf32> return %0: vector<1xf32> } + +// ----- + +// CHECK-LABEL: transfer_write_scalable +func.func @transfer_write_scalable(%arg0: memref>, %arg1: f32) { + %0 = llvm.mlir.constant(0 : i32) : i32 + %c0 = arith.constant 0 : index + %dim = memref.dim %arg0, %c0 : memref> + %1 = llvm.intr.experimental.stepvector : vector<[16]xi32> + %2 = arith.index_cast %dim : index to i32 + %3 = llvm.mlir.undef : vector<[16]xi32> + %4 = llvm.insertelement %2, %3[%0 : i32] : vector<[16]xi32> + %5 = llvm.shufflevector %4, %3 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] : vector<[16]xi32> + %6 = arith.cmpi slt, %1, %5 : vector<[16]xi32> + %7 = llvm.mlir.undef : vector<[16]xf32> + %8 = llvm.insertelement %arg1, %7[%0 : i32] : vector<[16]xf32> + %9 = llvm.shufflevector %8, %7 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] : vector<[16]xf32> + vector.transfer_write %9, %arg0[%c0], %6 {in_bounds = [true]} : vector<[16]xf32>, memref> + return +} + +// CHECK-SAME: %[[ARG_0:.*]]: memref>, +// CHECK: %[[C_0:.*]] = arith.constant 0 : index +// CHECK: %[[C_16:.*]] = arith.constant 16 : index +// CHECK: %[[STEP:.*]] = arith.constant 1 : index +// CHECK: %[[MASK_VEC:.*]] = arith.cmpi slt, %{{.*}}, %{{.*}} : vector<[16]xi32> +// CHECK: %[[VSCALE:.*]] = vector.vscale +// CHECK: %[[UB:.*]] = arith.muli %[[VSCALE]], %[[C_16]] : index +// CHECK: scf.for %[[IDX:.*]] = %[[C_0]] to %[[UB]] step %[[STEP]] { +// CHECK: %[[MASK_VAL:.*]] = vector.extractelement %[[MASK_VEC]][%[[IDX]] : index] : vector<[16]xi1> +// CHECK: scf.if %[[MASK_VAL]] { +// CHECK: %[[VAL_TO_STORE:.*]] = vector.extractelement %{{.*}}[%[[IDX]] : index] : vector<[16]xf32> +// CHECK: memref.store %[[VAL_TO_STORE]], %[[ARG_0]][%[[IDX]]] : memref> +// CHECK: } else { +// CHECK: } +// CHECK: }