2929#include " mlir/Support/LLVM.h"
3030#include " llvm/ADT/TypeSwitch.h"
3131#include " llvm/Support/Debug.h"
32+ #include " llvm/Support/DebugLog.h"
3233#include " llvm/Support/InterleavedRange.h"
3334#include " llvm/Support/raw_ostream.h"
3435#include < utility>
3839using namespace mlir ;
3940using namespace mlir ::linalg;
4041
41- #define DBGS () (llvm::dbgs() << " [" DEBUG_TYPE << " ]: " )
42- #define DBGSNL () (llvm::dbgs() << " \n " )
43-
4442// ===----------------------------------------------------------------------===//
4543// Transformations exposed as functional-style API calls.
4644// ===----------------------------------------------------------------------===//
@@ -276,23 +274,18 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
276274 tensor::PadOp::create (rewriter, loc, collapsed, packOp.getSource (), lows,
277275 highs, paddingValue, /* nofold=*/ false );
278276
279- LLVM_DEBUG (
280- DBGSNL (); DBGSNL ();
281- DBGS () << " insertPositions: "
282- << llvm::interleaved (packingMetadata.insertPositions );
283- DBGSNL (); DBGS () << " outerPositions: "
284- << llvm::interleaved (packingMetadata.outerPositions );
285- DBGSNL (); DBGS () << " packedShape: "
286- << llvm::interleaved (packedTensorType.getShape ());
287- DBGSNL (); DBGS () << " packedToStripMinedShapePerm: "
288- << llvm::interleaved (packedToStripMinedShapePerm);
289- DBGSNL ();
290- DBGS () << " reassociations: "
291- << llvm::interleaved (llvm::map_range (
292- packingMetadata.reassociations , stringifyReassocIndices));
293- DBGSNL ();
294- DBGS () << " stripMinedShape: " << llvm::interleaved (stripMinedShape);
295- DBGSNL (); DBGS () << " collapsed type: " << collapsed; DBGSNL (););
277+ LDBG () << " insertPositions: "
278+ << llvm::interleaved (packingMetadata.insertPositions );
279+ LDBG () << " outerPositions: "
280+ << llvm::interleaved (packingMetadata.outerPositions );
281+ LDBG () << " packedShape: " << llvm::interleaved (packedTensorType.getShape ());
282+ LDBG () << " packedToStripMinedShapePerm: "
283+ << llvm::interleaved (packedToStripMinedShapePerm);
284+ LDBG () << " reassociations: "
285+ << llvm::interleaved (llvm::map_range (packingMetadata.reassociations ,
286+ stringifyReassocIndices));
287+ LDBG () << " stripMinedShape: " << llvm::interleaved (stripMinedShape);
288+ LDBG () << " collapsed type: " << collapsed;
296289
297290 if (lowerPadLikeWithInsertSlice && packOp.isLikePad ()) {
298291 // Pack ops which operate as simple pads may not produce legal
@@ -317,7 +310,7 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
317310 rewriter, loc, /* source=*/ padOp, /* dest=*/ packOp.getDest (),
318311 /* offsets=*/ zeros, sizes, /* strides=*/ ones);
319312
320- LLVM_DEBUG ( DBGS ( ) << " insert_slice op: " << insertSliceOp; DBGSNL ();) ;
313+ LDBG ( ) << " insert_slice op: " << insertSliceOp;
321314
322315 rewriter.replaceOp (packOp, insertSliceOp->getResults ());
323316
@@ -339,10 +332,9 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
339332 auto transposeOp = linalg::TransposeOp::create (
340333 rewriter, loc, reshapeOp.getResult (), packOp.getDest (), transpPerm);
341334
342- LLVM_DEBUG (DBGSNL (); DBGSNL (); DBGSNL ();
343- DBGS () << " reshape op: " << reshapeOp; DBGSNL ();
344- DBGS () << " transpPerm: " << llvm::interleaved (transpPerm);
345- DBGSNL (); DBGS () << " transpose op: " << transposeOp; DBGSNL (););
335+ LDBG () << " reshape op: " << reshapeOp;
336+ LDBG () << " transpPerm: " << llvm::interleaved (transpPerm);
337+ LDBG () << " transpose op: " << transposeOp;
346338
347339 // 7. Replace packOp by transposeOp.
348340 rewriter.replaceOp (packOp, transposeOp->getResults ());
@@ -410,21 +402,16 @@ linalg::lowerUnPack(RewriterBase &rewriter, linalg::UnPackOp unPackOp,
410402 linalg::TransposeOp::create (rewriter, loc, unPackOp.getSource (), emptyOp,
411403 packedToStripMinedShapePerm);
412404
413- LLVM_DEBUG (
414- DBGSNL (); DBGSNL ();
415- DBGS () << " insertPositions: "
416- << llvm::interleaved (packingMetadata.insertPositions );
417- DBGSNL (); DBGS () << " packedShape: "
418- << llvm::interleaved (packedTensorType.getShape ());
419- DBGSNL (); DBGS () << " packedToStripMinedShapePerm: "
420- << llvm::interleaved (packedToStripMinedShapePerm);
421- DBGSNL ();
422- DBGS () << " reassociations: "
423- << llvm::interleaved (llvm::map_range (
424- packingMetadata.reassociations , stringifyReassocIndices));
425- DBGSNL ();
426- DBGS () << " stripMinedShape: " << llvm::interleaved (stripMinedShape);
427- DBGSNL (); DBGS () << " collapsed type: " << collapsedType; DBGSNL (););
405+ LDBG () << " insertPositions: "
406+ << llvm::interleaved (packingMetadata.insertPositions );
407+ LDBG () << " packedShape: " << llvm::interleaved (packedTensorType.getShape ());
408+ LDBG () << " packedToStripMinedShapePerm: "
409+ << llvm::interleaved (packedToStripMinedShapePerm);
410+ LDBG () << " reassociations: "
411+ << llvm::interleaved (llvm::map_range (packingMetadata.reassociations ,
412+ stringifyReassocIndices));
413+ LDBG () << " stripMinedShape: " << llvm::interleaved (stripMinedShape);
414+ LDBG () << " collapsed type: " << collapsedType;
428415
429416 // 4. Collapse from the stripMinedShape to the padded result.
430417 auto reshapeOp = tensor::CollapseShapeOp::create (
@@ -486,10 +473,9 @@ FailureOr<PackResult> linalg::pack(RewriterBase &rewriter,
486473 SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray ();
487474 SmallVector<utils::IteratorType> iteratorTypes =
488475 linalgOp.getIteratorTypesArray ();
489- LLVM_DEBUG (DBGS () << " Start packing: " << linalgOp << " \n "
490- << " maps: " << llvm::interleaved (indexingMaps) << " \n "
491- << " iterators: " << llvm::interleaved (iteratorTypes)
492- << " \n " );
476+ LDBG () << " Start packing: " << linalgOp;
477+ LDBG () << " maps: " << llvm::interleaved (indexingMaps);
478+ LDBG () << " iterators: " << llvm::interleaved (iteratorTypes);
493479
494480 SmallVector<linalg::PackOp> packOps;
495481 SmallVector<linalg::UnPackOp> unPackOps;
@@ -511,14 +497,11 @@ FailureOr<PackResult> linalg::pack(RewriterBase &rewriter,
511497 packedOperandsDims.packedDimForEachOperand = *maybePackedDimForEachOperand;
512498 listOfPackedOperandsDim.pushBack (std::move (packedOperandsDims));
513499
514- LLVM_DEBUG (
515- DBGS () << " ++++ After pack size #" << i << " : " << packedSizes[i]
516- << " \n "
517- << " maps: " << llvm::interleaved (indexingMaps) << " \n "
518- << " iterators: " << llvm::interleaved (iteratorTypes) << " \n "
519- << " packedDimForEachOperand: "
520- << llvm::interleaved (packedOperandsDims.packedDimForEachOperand )
521- << " \n " );
500+ LDBG () << " ++++ After pack size #" << i << " : " << packedSizes[i];
501+ LDBG () << " maps: " << llvm::interleaved (indexingMaps);
502+ LDBG () << " iterators: " << llvm::interleaved (iteratorTypes);
503+ LDBG () << " packedDimForEachOperand: "
504+ << llvm::interleaved (packedOperandsDims.packedDimForEachOperand );
522505 }
523506
524507 // Step 2. Propagate packing to all LinalgOp operands.
@@ -534,10 +517,9 @@ FailureOr<PackResult> linalg::pack(RewriterBase &rewriter,
534517 listOfPackedOperandsDim.extractPackedDimsForOperand (pos);
535518 SmallVector<OpFoldResult> innerPackSizes =
536519 listOfPackedOperandsDim.extractPackSizesForOperand (pos);
537- LLVM_DEBUG (DBGS () << " operand: " << operand << " \n "
538- << " innerPos: " << llvm::interleaved (innerPos) << " \n "
539- << " innerPackSizes: "
540- << llvm::interleaved (innerPackSizes) << " \n " );
520+ LDBG () << " operand: " << operand;
521+ LDBG () << " innerPos: " << llvm::interleaved (innerPos);
522+ LDBG () << " innerPackSizes: " << llvm::interleaved (innerPackSizes);
541523 if (innerPackSizes.empty ()) {
542524 inputsAndInits.push_back (operand);
543525 continue ;
@@ -776,8 +758,8 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
776758
777759 int64_t numLoops = linalgOp.getNumLoops ();
778760 if (numLoops <= 2 ) {
779- LLVM_DEBUG ( DBGS ( ) << " need 3+ loops to find a matmul to pack, got "
780- << numLoops << " \n in : " << linalgOp << " \n " ) ;
761+ LDBG ( ) << " need 3+ loops to find a matmul to pack, got " << numLoops
762+ << " in : " << linalgOp;
781763 return rewriter.notifyMatchFailure (
782764 linalgOp, " need 3+ loops to find a matmul to pack" );
783765 }
@@ -801,8 +783,7 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
801783 FailureOr<ContractionDimensions> maybeDimensions =
802784 inferContractionDims (linalgOp);
803785 if (failed (maybeDimensions)) {
804- LLVM_DEBUG (DBGS () << " couldn't infer matmul iterators in: " << linalgOp
805- << " \n " );
786+ LDBG () << " couldn't infer matmul iterators in: " << linalgOp;
806787 return rewriter.notifyMatchFailure (linalgOp,
807788 " couldn't infer matmul iterators" );
808789 }
@@ -814,10 +795,8 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
814795 // to plug a heuristic.
815796 int64_t mPos = maybeDimensions->m .back (), nPos = maybeDimensions->n .back (),
816797 kPos = maybeDimensions->k .back ();
817- LLVM_DEBUG (DBGSNL (); DBGSNL (); DBGSNL ();
818- DBGS () << " Start packing generic op greedily with (m@" << mPos
819- << " , n@" << nPos << " , k@" << kPos << " ): " << linalgOp
820- << " \n " ;);
798+ LDBG () << " Start packing generic op greedily with (m@" << mPos << " , n@"
799+ << nPos << " , k@" << kPos << " ): " << linalgOp;
821800
822801 // 2.a. Rewrite as a generic.
823802 auto genericOp = dyn_cast<GenericOp>(linalgOp.getOperation ());
@@ -833,14 +812,14 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
833812 // not change the indexings of any operand.
834813 SmallVector<int64_t > permutation =
835814 computePermutationVector (numLoops, {mPos , nPos, kPos }, mmnnkkPos);
836- LLVM_DEBUG ( DBGS ( ) << " perm: " << llvm::interleaved (permutation) << " \n " );
815+ LDBG ( ) << " perm: " << llvm::interleaved (permutation);
837816 // Sign .. unsigned pollution.
838817 SmallVector<unsigned > unsignedPerm (permutation.begin (), permutation.end ());
839818 FailureOr<GenericOp> interchangeResult =
840819 interchangeGenericOp (rewriter, genericOp, unsignedPerm);
841820 assert (succeeded (interchangeResult) && " unexpected failure interchanging op" );
842821 genericOp = *interchangeResult;
843- LLVM_DEBUG ( DBGS ( ) << " Generalized Op to pack: " << genericOp << " \n " ;) ;
822+ LDBG ( ) << " Generalized Op to pack: " << genericOp;
844823
845824 // At this point, the op iterators are normalized to {leading, k, m, n}.
846825 // The layouts induced by packing will always be:
@@ -862,12 +841,11 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
862841
863842 // Add leading zeros to match numLoops, we only pack the last 3 dimensions
864843 // post interchange.
865- LLVM_DEBUG (DBGS () << " paddedSizesNextMultipleOf: "
866- << llvm::interleaved (paddedSizesNextMultipleOf) << " \n "
867- << " loopRanges: "
868- << llvm::interleaved (llvm::map_range (
869- loopRanges, [](Range r) { return r.size ; }))
870- << " \n " );
844+ LDBG () << " paddedSizesNextMultipleOf: "
845+ << llvm::interleaved (paddedSizesNextMultipleOf);
846+ LDBG () << " loopRanges: "
847+ << llvm::interleaved (
848+ llvm::map_range (loopRanges, [](Range r) { return r.size ; }));
871849 SmallVector<OpFoldResult> adjustedPackedSizes (numLoops - packedSizes.size (),
872850 rewriter.getIndexAttr (0 ));
873851 for (int64_t i = 0 , e = numPackedDims; i < e; ++i) {
@@ -883,8 +861,7 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
883861 {loopRanges[adjustedPackedSizes.size ()].size ,
884862 rewriter.getIndexAttr (paddedSizesNextMultipleOf[i])}));
885863 }
886- LLVM_DEBUG (DBGS () << " adjustedPackedSizes: "
887- << llvm::interleaved (adjustedPackedSizes) << " \n " );
864+ LDBG () << " adjustedPackedSizes: " << llvm::interleaved (adjustedPackedSizes);
888865
889866 // TODO: If we wanted to give the genericOp a name after packing, after
890867 // calling `pack` would be a good time. One would still need to check that
@@ -1214,9 +1191,8 @@ LogicalResult DecomposeOuterUnitDimsPackOpPattern::matchAndRewrite(
12141191 }
12151192 srcPermForTranspose.append (innerDimPos.begin (), innerDimPos.end ());
12161193
1217- LLVM_DEBUG (DBGS () << " Pack permutation: " << packOp << " \n "
1218- << " perm: " << llvm::interleaved (srcPermForTranspose)
1219- << " \n " );
1194+ LDBG () << " Pack permutation: " << packOp;
1195+ LDBG () << " perm: " << llvm::interleaved (srcPermForTranspose);
12201196
12211197 // 2.1 Create tensor.empty (init value for TransposeOp)
12221198 SmallVector<OpFoldResult> transShapeForEmptyOp (srcRank - numTiles,
0 commit comments