[Mlir-commits] [mlir] Rewrites for I2 to I8 signed and unsigned extension (PR #121298)

[Han-Chung Wang]

================
@@ -1172,70 +1172,167 @@ Value BitCastRewriter::genericRewriteStep(
   return runningResult;
 }
 
+Value bitcastSubByteVectorToI8(PatternRewriter &rewriter, Location loc,
+                               Value srcValue) {
+  VectorType srcVecType = cast<VectorType>(srcValue.getType());
+  int64_t srcBitwidth = srcVecType.getElementType().getIntOrFloatBitWidth();
+  assert(srcBitwidth % 8 != 0 && "Invalid source bitwidth");
+  int64_t bitwidthFactor = 8 / srcBitwidth;
+  SmallVector<int64_t> i8VecShape = llvm::to_vector(srcVecType.getShape());
+  i8VecShape.back() = i8VecShape.back() / bitwidthFactor;
+  auto i8VecType = VectorType::get(i8VecShape, rewriter.getI8Type());
+  return rewriter.create<vector::BitCastOp>(loc, i8VecType, srcValue);
+}
----------------
hanhanW wrote:

Few suggestions:

- Add a documentation to the function.
- Add `static` keyword to the local function. Same for the below function. https://llvm.org/docs/CodingStandards.html#anonymous-namespaces
- Use auto for `srcVecType` because `cast` already spells the type. https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable
- [optional] I'd rename `i8VecShape` to `vecShape`.
- You do not need the llvm::to_vector. One of the SmallVector constructors takes ArrayRef type directly. So it can be `SmallVector<int64_t> vecShape(srcVecType.getShape())`.

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