b/376639863: Speed up int4 unpacking by using inline_asm with packed elements at
an early stage on the XLA side with further optimizations from nVidia #24 and
#25 on the Triton side.

--- a/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
+++ b/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp
@@ -198,6 +198,19 @@ struct JoinOpConversion : public ConvertOpToLLVMPattern<JoinOp> {
     assert(lhsVals.size() == rhsVals.size());
     SmallVector<Value> joinedVals;
     joinedVals.resize(lhsVals.size() * 2);
+
+    // Specifically for packed upcasting from 4b to 16b dtypes
+    // numContiguousValues cannot be too large, since the two outputs of
+    // inline_asm contain interleaved values OTOH, if numContiguousValues * 16b
+    // < 32b, then we'll need to rearrange 16b values in 32b registers. Hence we
+    // set numContiguousValues to 2
+    auto inlineOp =
+        dyn_cast_or_null<ElementwiseInlineAsmOp>(op.getLhs().getDefiningOp());
+    if (inlineOp && inlineOp.getPackedElement() == 4 &&
+        dstTy.getElementTypeBitWidth() == 16 && joinedVals.size() > 2) {
+      numContiguousValues = 2;
+    }
+
     for (int i = 0; i < lhsVals.size(); i += numContiguousValues) {
       for (int j = 0; j < numContiguousValues; j++) {
         joinedVals[2 * i + j] = lhsVals[i + j];

--- a/lib/Dialect/TritonGPU/Transforms/Utility.cpp
+++ b/lib/Dialect/TritonGPU/Transforms/Utility.cpp
@@ -1140,6 +1140,39 @@ swizzleDotOperandLike(RankedTensorType type, ttg::CGAEncodingAttr cgaLayout) {
       type.getElementTypeBitWidth(), false);
 }
 
+// Rough utility for obtaining a SharedEnc for a LinearEncoding,
+// as we've replaced DotOpEnc with Linear in some cases
+// (specifically, fp4ToFp and similar unpack-upcast thru join)
+std::optional<ttg::SwizzledSharedEncodingAttr>
+getSharedForLinear(ttg::LinearEncodingAttr enc,
+                   ArrayRef<unsigned int> globalOrder, ArrayRef<int64_t> shape,
+                   unsigned elemBitWidth, ttg::CGAEncodingAttr cgaLayout) {
+  auto ctx = enc.getContext();
+  auto ll = enc.getLinearLayout();
+  auto rank = shape.size();
+  if (rank != 2)
+    return std::nullopt;
+  auto order = enc.getOrder();
+  assert(globalOrder.size() == rank);
+  // TODO add memdesc_trans support for dot(trans(cvt(src) #linear) #dot_op)
+  if (order != globalOrder)
+    return std::nullopt;
+  auto innerDim = order[0];
+  auto outerDim = order[1];
+  auto contigPerWarp = enc.getContigPerWarp();
+  constexpr unsigned BANK_SIZE{128};
+  auto elemBytes = elemBitWidth / 8;
+  auto vec = contigPerWarp[innerDim];
+  auto rowSize = elemBytes * (unsigned)shape[innerDim];
+  auto perPhase = std::max(BANK_SIZE / rowSize, 1u);
+  auto maxPhase = std::max(contigPerWarp[outerDim] / perPhase, 1u);
+  // cp.async does not support transfer size < 4B
+  if (vec * elemBytes < 4 && perPhase < maxPhase)
+    return std::nullopt;
+  return ttg::SwizzledSharedEncodingAttr::get(ctx, vec, perPhase, maxPhase,
+                                              order, cgaLayout);
+}
+
 // If all the transitive uses of the given value have are used by a convert to
 // the same dot operand encoding, return the shared encoding that needs to be
 // used to be compatible with users' layouts. If there are incompatible shared
@@ -1174,14 +1207,21 @@ getSharedEncIfAllUsersAreDotEnc(Value val, bool &incompatible) {
       auto CGALayout = isa<ttg::LinearEncodingAttr>(dstTy.getEncoding())
                            ? ttg::getCGALayout(srcTy.getEncoding())
                            : ttg::getCGALayout(dstTy.getEncoding());
+      auto order = getOrderForMemory(srcTy);
+      unsigned bitWidth = srcTy.getElementTypeBitWidth();
 
       if (auto dot =
               dyn_cast<ttg::DotOperandEncodingAttr>(dstTy.getEncoding())) {
-        auto order = getOrderForMemory(srcTy);
-        unsigned bitWidth = srcTy.getElementTypeBitWidth();
         tempAttr = ttg::SwizzledSharedEncodingAttr::get(
             val.getContext(), dot, srcTy.getShape(), order, CGALayout, bitWidth,
             /*needTrans=*/false);
+      } else if (auto linearEnc =
+                     dyn_cast<ttg::LinearEncodingAttr>(dstTy.getEncoding())) {
+        auto attrOpt = getSharedForLinear(linearEnc, order, srcTy.getShape(),
+                                          bitWidth, CGALayout);
+        if (!attrOpt)
+          return std::nullopt;
+        tempAttr = *attrOpt;
       } else {
         // Try to see if the layout is like an mma microtile
         tempAttr = swizzleDotOperandLike(dstTy, CGALayout);