Cleanup

Author: leonling-ll

python/src/ir.cc

@@ -1622,7 +1622,7 @@ void init_triton_ir(py::module &&m) {
              if (haveDump) {
                auto printingFlags = OpPrintingFlags();
                printingFlags.elideLargeElementsAttrs(16);
-               //  printingFlags.enableDebugInfo();
+               printingFlags.enableDebugInfo();
                auto printAlways = [funcToDump](Pass *, Operation *op) -> bool {
                  if (funcToDump.empty())
                    return true;

third_party/intel/lib/Dialect/TritonIntelGPU/IR/LinearLayoutConversions.cpp

@@ -56,8 +56,6 @@ LinearLayout identityND(StringAttr inDimName, ArrayRef<unsigned> shape,
   LinearLayout ret = LinearLayout::empty();
   for (int i = 0; i < shape.size(); i++) {
     // Start with the most-minor dimension, which is order[0].
-    std::cout << "i: " << i << " shape[i]: " << shape[i]
-              << " order[i]: " << order[i] << std::endl;
     int dim = order[i];
     ret *= LinearLayout::identity1D(shape[dim], inDimName, outDimNames[dim]);
   }
@@ -291,15 +289,6 @@ LinearLayout ensureLayoutNotSmallerThan(
     assert(actualSize > desiredSize ||
            desiredSize % actualSize == 0 && "bad shape");
     ret *= LinearLayout::identity1D(desiredSize / actualSize, kDim, outDimName);
-    std::cout << "actualSize: " << actualSize << " desiredSize: " << desiredSize
-              << std::endl;
-    std::cout << "outDimName: " << outDimName.str() << std::endl;
-    std::cout << "identity1D: "
-              << LinearLayout::identity1D(desiredSize / actualSize, kDim,
-                                          outDimName)
-                     .toString()
-              << std::endl;
-    std::cout << "ret: " << ret.toString() << std::endl;
     assert(ret.getOutDimSize(outDimName) >= desiredSize && "bad grow");
   }
   return ret;
@@ -323,12 +312,6 @@ LinearLayout combineCtaCgaWithShape(LinearLayout ctaLayout,
 
   SmallVector<StringAttr> outDimNames = standardOutDimNames(ctx, rank);
 
-  std::cout << "shape: ";
-  for (auto s : shape) {
-    std::cout << s << ", ";
-  }
-
-  std::cout << std::endl;
   llvm::SmallDenseMap<StringAttr, int64_t> labeledShape;
   for (auto [dim, size] : llvm::zip(outDimNames, shape)) {
     labeledShape[dim] = size;
@@ -337,7 +320,6 @@ LinearLayout combineCtaCgaWithShape(LinearLayout ctaLayout,
   LinearLayout cgaLayout =
       ensureLayoutNotLargerThan(makeCgaLayout(cgaLayoutAttr), labeledShape)
           .transposeOuts(llvm::to_vector(ctaLayout.getOutDimNames()));
-  std::cout << "\ncgaLayout: " << cgaLayout.toString() << std::endl;
 
   // Calculate the shape of the ctaLayout, which is `shape` divided by the
   // cgaLayout's size.
@@ -346,29 +328,19 @@ LinearLayout combineCtaCgaWithShape(LinearLayout ctaLayout,
              llvm::to_vector(cgaLayout.getOutDimNames()) &&
          "bad layout");
 
-  std::cout << "ctaShape: ";
   for (auto dim : ctaLayout.getOutDimNames()) {
     ctaShape[dim] =
         std::max(int64_t{1}, labeledShape[dim] / cgaLayout.getOutDimSize(dim));
-    std::cout << ctaShape[dim] << ", ";
   }
-  std::cout << std::endl;
 
   ctaLayout = ensureLayoutNotSmallerThan(ctaLayout, ctaShape);
-  std::cout << "\nctaLayout not smaller than: " << ctaLayout.toString()
-            << std::endl;
   ctaLayout = ensureLayoutNotLargerThan(ctaLayout, ctaShape);
-  std::cout << "\nctaLayout not larger than: " << ctaLayout.toString()
-            << std::endl;
 
-  std::cout << "\ncta * cga: " << (ctaLayout * cgaLayout).toString()
-            << std::endl;
   LinearLayout ret =
       (std::move(ctaLayout) * std::move(cgaLayout)).transposeOuts(outDimNames);
   for (auto dim : ret.getOutDimNames()) {
     assert(ret.getOutDimSize(dim) == labeledShape[dim] && "bad shape");
   }
-  std::cout << "\ncombineCtaCgaWithShape: " << ret.toString() << std::endl;
   return ret;
 }
 
@@ -565,7 +537,6 @@ LinearLayout DPAStoLinearLayout(ArrayRef<int64_t> shape, Attribute layout,
           LinearLayout::identity1D(warpsPerCTA[0], kWarp, outDimNames[0]);
 
   } else if (opIdx == 1) { // Operand B
-    std::cout << "\nOperand B" << std::endl;
     auto regBasesB = DPASRegBasesB(opsPerChannel, executionSize, threadsPerWarp,
                                    systolicDepth);
     auto laneBasesB =
@@ -587,32 +558,20 @@ LinearLayout DPAStoLinearLayout(ArrayRef<int64_t> shape, Attribute layout,
       tileLayout *=
           LinearLayout::identity1D(warpsPerCTA[0], kWarp, outDimNames[0]);
   } else { // opIdx=2 -> Operand C
-    std::cout << "\nOperand C" << std::endl;
     auto regBasesC = DPASRegBasesC(repeatCount, executionSize, threadsPerWarp);
     auto laneBasesC =
         DPASLaneBasesC(repeatCount, executionSize, threadsPerWarp);
     tileLayout = LinearLayout({{kRegister, regBasesC}, {kLane, laneBasesC}},
                               ArrayRef(outDimNames).take_back(2));
-    std::cout << tileLayout.toString() << std::endl;
     // The per-inst layout is repeated at each repCluster.
     // Hence, multiply with the identity layouts starting from the
     // least significant dimension.
     dimNonK = rank - 2;
     dimK = rank - 1;
     tileLayout *= LinearLayout::identity1D(repCluster[dimK], kRegister,
                                            outDimNames[dimK]);
-    std::cout << (LinearLayout::identity1D(repCluster[dimK], kRegister,
-                                           outDimNames[dimK])
-                      .toString())
-              << std::endl;
-    std::cout << (tileLayout.toString()) << std::endl;
     tileLayout *= LinearLayout::identity1D(repCluster[dimNonK], kRegister,
                                            outDimNames[dimNonK]);
-    std::cout << (LinearLayout::identity1D(repCluster[dimNonK], kRegister,
-                                           outDimNames[dimNonK])
-                      .toString())
-              << std::endl;
-    std::cout << (tileLayout.toString()) << std::endl;
 
     // // The identical layout is repeated among warps
     tileLayout *=
@@ -622,23 +581,8 @@ LinearLayout DPAStoLinearLayout(ArrayRef<int64_t> shape, Attribute layout,
     if (rank == 3)
       tileLayout *=
           LinearLayout::identity1D(warpsPerCTA[0], kWarp, outDimNames[0]);
-    // std::cout << (tileLayout.toString()) << std::endl;
   }
 
-  // Lastly, the layout repeats to match the shape.
-  // Operand A/B repeats through the K-dimension first then repeats
-  // through the non-K dimension.
-  // SmallVector<int64_t> numReps = dpas.getDPASRepetitions(shape, opIdx);
-  // std::cout << "numReps: " << numReps[0] << ", " << numReps[1] << std::endl;
-  // tileLayout *=
-  //     LinearLayout::identity1D(numReps[dimK], kRegister, outDimNames[dimK]);
-  // tileLayout *= LinearLayout::identity1D(numReps[dimNonK], kRegister,
-  //                                        outDimNames[dimNonK]);
-  // if (rank == 3)
-  //   tileLayout *=
-  //       LinearLayout::identity1D(numReps[0], kRegister, outDimNames[0]);
-  // std::cout << (tileLayout.toString()) << std::endl;
-
   return combineCtaCgaWithShape(std::move(tileLayout),
                                 CTALayoutAttr::getDefault(ctx, rank), shape);
 }