hack me

Signed-off-by: dchigarev <[email protected]>

Author: dchigarev

lib/Conversion/TritonGPUToLLVM/Utility.cpp

@@ -367,12 +367,13 @@ Value getSmemVecAddrNEW(const LinearLayout &regLayout,
   // solution for all swizzled shared memory scenarios, including the edge case
   // mentioned above.
   if (isSimpleSharedMemoryAccess(shape, allocShape, sharedEnc)) { // Case 1
-    smemOffset = applyLinearLayout(loc, rewriter, regToSharedLayout,
+    auto res = applyLinearLayout(loc, rewriter, regToSharedLayout,
                                    {{kRegister, regId},
                                     {kLane, laneId},
                                     {kWarp, warpId},
-                                    {kBlock, blockId}})[0]
-                     .second;
+                                    {kBlock, blockId}});
+    std::cout << "linearLayRes.size(): " << res.size() << "\n";
+    smemOffset = res[0].second;
   } else { // Case 2 -> rank-reduced swizzling
     assert(rank >= 2 && "Swizzling only applies to tensors with rank >= 2");
     assert(!sharedEnc.getHasLeadingOffset() &&
@@ -426,7 +427,7 @@ Value getSmemVecAddrNEW(const LinearLayout &regLayout,
 } // namespace
 
 
-bool getBoolFromEnv(const std::string& envVar, bool defaultValue = false) {
+static bool getBoolFromEnv(const std::string& envVar, bool defaultValue = false) {
     const char* value = std::getenv(envVar.c_str());
     if (value == nullptr) {
         return defaultValue; // Return default if the variable is not set
@@ -549,10 +550,18 @@ bool emitTransferBetweenRegistersAndSharedNEW(
   StringAttr kWarp = str_attr("warp");
 
   auto shape = sharedTy.getShape();
+  llvm::dbgs() << "registerTy enc\n";
+  registerTy.dump();
+  registerTy.getEncoding().dump();
+  llvm::dbgs() << "shape: "; for (auto &el : shape) { llvm::dbgs() << el << " ";} llvm::dbgs() << "\n";
   LinearLayout regLayout =
       triton::gpu::toLinearLayout(shape, registerTy.getEncoding());
   printLinearThing(regLayout, "regLayout");
 
+  llvm::dbgs() << "sharedTy enc\n";
+  sharedTy.dump();
+  sharedTy.getEncoding().dump();
+  llvm::dbgs() << "shape: "; for (auto &el : shape) { llvm::dbgs() << el << " ";} llvm::dbgs() << "\n";
   LinearLayout sharedLayout = triton::gpu::toLinearLayout(
       shape, sharedTy.getEncoding(), elemLlvmTy.getIntOrFloatBitWidth());
   printLinearThing(sharedLayout, "sharedLayout");
@@ -653,13 +662,30 @@ SmallVector<Value> loadSharedToDistributed(RankedTensorType dstTy,
   bool success = emitTransferBetweenRegistersAndShared(
       dstTy, srcTy, elemLlvmTy, /*maxVecElems=*/std::nullopt, smemObj, loc,
       rewriter, target, [&](VectorType vecTy, Value vecAddr) {
-        auto vecVal = load(vecTy, vecAddr);
-        vecVal.setAlignment(vecTy.getNumElements() *
-                            elemLlvmTy.getIntOrFloatBitWidth() / 8);
-
-        for (int v = 0; v < vecTy.getNumElements(); v++) {
-          ret.push_back(extract_element(elemLlvmTy, vecVal, i32_val(v)));
+        if (vecTy.getNumElements() >= 64) {
+            assert(vecTy.getNumElements() % 64 == 0);
+            for (int i = 0; i < vecTy.getNumElements(); i+=64) {
+                auto smallVecTy = vec_ty(elemLlvmTy, 64);
+                auto vecAddrNew = gep(vecAddr.getType(), i32_ty, vecAddr, SmallVector<Value>({i32_val(i)}));
+                auto vecVal = load(smallVecTy, vecAddrNew);
+                vecVal.setAlignment(smallVecTy.getNumElements() *
+                                    elemLlvmTy.getIntOrFloatBitWidth() / 8);
+
+                for (int v = 0; v < 64; v++) {
+                    ret.push_back(extract_element(elemLlvmTy, vecVal, i32_val(v)));
+                }
+            }
+            
+        } else {
+            auto vecVal = load(vecTy, vecAddr);
+            vecVal.setAlignment(vecTy.getNumElements() *
+                                elemLlvmTy.getIntOrFloatBitWidth() / 8);
+
+            for (int v = 0; v < vecTy.getNumElements(); v++) {
+                ret.push_back(extract_element(elemLlvmTy, vecVal, i32_val(v)));
+            }
         }
+
       });
   if (!success)
     llvm::report_fatal_error("Failed to emit transfer from shared to register");

python/triton/runtime/build.py

@@ -93,5 +93,13 @@ def _build(name, src, srcdir, library_dirs, include_dirs, libraries, extra_compi
     if os.getenv("VERBOSE"):
         print(" ".join(cc_cmd))
 
-    subprocess.check_call(cc_cmd, stdout=subprocess.DEVNULL)
+    result = subprocess.run(cc_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
+
+    if result.returncode != 0:
+        print(f"Error: Command failed with exit code {result.returncode}")
+        if result.stderr:
+            print("Error output:", result.stderr)
+
+    # breakpoint()
+    # subprocess.check_call(cc_cmd, stdout=subprocess.DEVNULL)
     return so

third_party/intel/lib/TritonIntelGPUToLLVM/MemoryOpToLLVM.cpp

@@ -20,6 +20,18 @@ using namespace mlir;
 using namespace mlir::triton;
 using namespace mlir::triton::gpu;
 
+static bool getBoolFromEnv(const std::string& envVar, bool defaultValue = false) {
+    const char* value = std::getenv(envVar.c_str());
+    if (value == nullptr) {
+        return defaultValue; // Return default if the variable is not set
+    }
+
+    std::string strValue(value);
+    for (char& c : strValue) c = std::tolower(c); // Convert to lowercase
+
+    return (strValue == "1" || strValue == "true" || strValue == "yes" || strValue == "on");
+}
+
 // blocked -> shared.
 // Swizzling in shared memory to avoid bank conflict. Normally used for
 // A/B operands of dots.
@@ -78,11 +90,14 @@ struct LocalAllocOpConversion
   LogicalResult
   matchAndRewrite(triton::gpu::LocalAllocOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    llvm::dbgs() << "\n\n===LocalAllocOpConversion BEFORE===\n";
-    op.dump();
-    for (auto& x : *(op->getParentRegion())) {
-        x.dump();
+    if (getBoolFromEnv("TR_LONG_IR")) {
+        llvm::dbgs() << "\n\n===LocalAllocOpConversion BEFORE===\n";
+        op.dump();
+        for (auto& x : *(op->getParentRegion())) {
+            x.dump();
+        }
     }
+
     if (!op.isSharedMemoryAlloc())
       return failure();
     Location loc = op->getLoc();
@@ -99,17 +114,24 @@ struct LocalAllocOpConversion
                                       loc, rewriter);
     // If there is an initial tensor, store it into the shared memory.
     if (op.getSrc()) {
+      llvm::dbgs() << "LocalAllocOp adaptor.src():\n";
+      adaptor.getSrc().dump();
+      llvm::dbgs() << "LocalAllocOp op.src() and op itself:\n";
+      op.getSrc().dump();
+      op.dump();
       lowerDistributedToShared(loc, op.getSrc(), op.getResult(),
                                adaptor.getSrc(), smemObj, typeConverter,
                                rewriter, targetInfo);
     }
     auto retVal = getStructFromSharedMemoryObject(loc, smemObj, rewriter);
     rewriter.replaceOp(op, retVal);
+    if (getBoolFromEnv("TR_LONG_IR")) {
     llvm::dbgs() << "\n\n===LocalAllocOpConversion AFTER===\n";
     retVal.dump();
     for (auto& x : *(retVal.getParentRegion())) {
         x.dump();
     }
+    }
     return success();
   }
 
@@ -195,11 +217,13 @@ struct LocalLoadOpConversion : public ConvertOpToLLVMPattern<LocalLoadOp> {
   lowerSharedToDistributed(LocalLoadOp op, LocalLoadOpAdaptor adaptor,
                            const LLVMTypeConverter *typeConverter,
                            ConversionPatternRewriter &rewriter) const {
+    if (getBoolFromEnv("TR_LONG_IR")) {
     llvm::dbgs() << "\n\n===lowerSharedToDistributed BEFORE===\n";
     op.dump();
     for (auto& x : *(op->getParentRegion())) {
         x.dump();
     }
+    }
     auto loc = op.getLoc();
     auto srcTy = op.getSrc().getType();
     auto dstTy = op.getResult().getType();
@@ -214,11 +238,13 @@ struct LocalLoadOpConversion : public ConvertOpToLLVMPattern<LocalLoadOp> {
 
     Value result = packLLElements(loc, typeConverter, outVals, rewriter, dstTy);
     rewriter.replaceOp(op, result);
+    if (getBoolFromEnv("TR_LONG_IR")) {
     llvm::dbgs() << "\n\n===lowerSharedToDistributed AFTER===\n";
     result.dump();
     for (auto& x : *(result.getParentRegion())) {
         x.dump();
     }
+    }
     return success();
   }