[Mlir-commits] [mlir] [mlir][Vectorizer] Added support to Vectorize tensor.unpack (PR #76087)

Diego Caballero llvmlistbot at llvm.org
Tue Feb 13 11:30:29 PST 2024


================
@@ -1420,16 +1420,28 @@ static Value createReadOrMaskedRead(OpBuilder &builder, Location loc,
   auto sourceShape = dyn_cast<ShapedType>(source.getType()).getShape();
   assert(sourceShape.size() == readShape.size());
   auto maskType = VectorType::get(readShape, builder.getI1Type());
-  auto vectorType = VectorType::get(readShape, padValue.getType());
+  Type vecElemType = padValue != nullptr
+                         ? padValue.getType()
+                         : cast<ShapedType>(source.getType()).getElementType();
+  auto vectorType = VectorType::get(readShape, vecElemType);
   int64_t readRank = readShape.size();
   auto zero = builder.create<arith::ConstantIndexOp>(loc, 0);
-  auto transferReadOp = builder.create<vector::TransferReadOp>(
-      loc,
-      /*vectorType=*/vectorType,
-      /*source=*/source,
-      /*indices=*/SmallVector<Value>(readRank, zero),
-      /*padding=*/padValue,
-      /*inBounds=*/SmallVector<bool>(readRank, true));
+  vector::TransferReadOp transferReadOp = nullptr;
+  if (padValue == nullptr) {
+    transferReadOp = builder.create<vector::TransferReadOp>(
+        loc,
+        /*vectorType=*/vectorType,
+        /*source=*/source,
+        /*indices=*/SmallVector<Value>(readRank, zero));
+  } else {
+    transferReadOp = builder.create<vector::TransferReadOp>(
+        loc,
+        /*vectorType=*/vectorType,
+        /*source=*/source,
+        /*indices=*/SmallVector<Value>(readRank, zero),
+        /*padding=*/padValue,
----------------
dcaballe wrote:

could you check if this argument is optional and we can pass `std::nullptr`? If so, we could remove the if/else above and have just one builder call.

https://github.com/llvm/llvm-project/pull/76087


More information about the Mlir-commits mailing list