skip mma

Author: liqiangxl

csrc/device_lower/analysis/tma.cpp

@@ -1180,8 +1180,13 @@ TMAInfo getTMAInfo(LoadStoreOp* ldst) {
 
   // Validate broadcast usage: TMA auto-fills out-of-bounds with zeros,
   // breaking broadcast semantics when broadcast dims participate in tile
-  // shape.
-  validateTMAConsumerBroadcasts(smem_tv);
+  // shape. Skip if smem_tv is a producer of mma op
+  const auto& consumers = ir_utils::consumerTvsOf(smem_tv);
+  if (std::none_of(consumers.begin(), consumers.end(), [](TensorView* tv) {
+        return tv->definition()->isA<MmaOp>();
+      })) {
+    validateTMAConsumerBroadcasts(smem_tv);
+  }
 
   MmaInputSmemSwizzle swizzle = getSwizzle(smem_tv);