Repository: llvm/llvm-project
File: mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp (C++)

Code before:
    return false;
  }

  auto isUnrollable = [](Value value, ArrayRef<int64_t> tileShape) {
    Type valTy = value.getType();
    if (auto tdescTy = dyn_cast<xegpu::TensorDescType>(valTy)) {
      xegpu::DistributeLayoutAttr layout = tdescTy.getLayoutAttr();
      return layout && !layout.getEffectiveInstDataAsInt().empty();
    }
    auto shapedType = dyn_cast<ShapedType>(valTy);
    return shapedType && !llvm::equal(tileShape, shapedType.getShape());
  };

  bool hasUnrollableOperands =
      llvm::any_of(op->getOpOperands(), [&](OpOperand &opr) {
        std::optional<SmallVector<int64_t>> tileShape = getTileShape(opr);
        return tileShape.has_value() && isUnrollable(opr.get(), *tileShape);
      });
  bool hasUnrollableResults =
      llvm::any_of(op->getOpResults(), [&](OpResult result) {
        std::optional<SmallVector<int64_t>> tileShape = getTileShape(result);
        return tileShape.has_value() && isUnrollable(result, *tileShape);
      });
  // ConvertLayoutOp must be processed to drop the inst_data in the layout
  bool isConvertLayoutWithInstData = false;
  if (isa<xegpu::ConvertLayoutOp>(op)) {
    xegpu::ConvertLayoutOp convertLayoutOp = cast<xegpu::ConvertLayoutOp>(op);
    auto targettLayout = convertLayoutOp.getTargetLayout();
    if (targettLayout && !targettLayout.getEffectiveInstDataAsInt().empty()) {
      isConvertLayoutWithInstData = true;
    }
  }
  return hasUnrollableOperands || hasUnrollableResults ||
         isConvertLayoutWithInstData;
}

void XeGPUBlockingPass::runOnOperation() {
  MLIRContext *ctx = &getContext();
  Operation *op = getOperation();

  if (!xegpu::recoverTemporaryLayouts(op)) {
    signalPassFailure();
    return;
  }

  auto getTileShapeAndCount = [](llvm::ArrayRef<int64_t> shape,
                                 xegpu::LayoutAttr layout) {
    int count = 1;
    SmallVector<int64_t> tileShape(shape);
    if (layout && layout.getInstData()) {
      DenseI32ArrayAttr instData = layout.getInstData();

Review comment (suggestion):
```suggestion
  if (auto convertLayoutOp = dyn_cast<xegpu::ConvertLayoutOp>(op)) {
```

Code after:
    return false;
  }

  auto isUnrollable = [](Value value, ArrayRef<int64_t> tileShape) {
    Type valTy = value.getType();
    if (auto tdescTy = dyn_cast<xegpu::TensorDescType>(valTy)) {
      xegpu::DistributeLayoutAttr layout = tdescTy.getLayoutAttr();
      return layout && !layout.getEffectiveInstDataAsInt().empty();
    }
    auto shapedType = dyn_cast<ShapedType>(valTy);
    return shapedType && !llvm::equal(tileShape, shapedType.getShape());
  };

  bool hasUnrollableOperands =
      llvm::any_of(op->getOpOperands(), [&](OpOperand &opr) {
        std::optional<SmallVector<int64_t>> tileShape = getTileShape(opr);
        return tileShape.has_value() && isUnrollable(opr.get(), *tileShape);
      });
  bool hasUnrollableResults =
      llvm::any_of(op->getOpResults(), [&](OpResult result) {
        std::optional<SmallVector<int64_t>> tileShape = getTileShape(result);
        return tileShape.has_value() && isUnrollable(result, *tileShape);
      });
  // ConvertLayoutOp must be processed to drop the inst_data in the layout
  bool isConvertLayoutWithInstData = false;
  if (auto convertLayoutOp = dyn_cast<xegpu::ConvertLayoutOp>(op)) {
    auto targettLayout = convertLayoutOp.getTargetLayout();
    if (targettLayout && !targettLayout.getEffectiveInstDataAsInt().empty()) {
      isConvertLayoutWithInstData = true;
    }
  }
  return hasUnrollableOperands || hasUnrollableResults ||
         isConvertLayoutWithInstData;
}

void XeGPUBlockingPass::runOnOperation() {
  MLIRContext *ctx = &getContext();
  Operation *op = getOperation();

  if (!xegpu::recoverTemporaryLayouts(op)) {
    signalPassFailure();
    return;
  }

  auto getTileShapeAndCount = [](llvm::ArrayRef<int64_t> shape,
                                 xegpu::LayoutAttr layout) {
    int count = 1;
    SmallVector<int64_t> tileShape(shape);
    if (layout && layout.getInstData()) {
      DenseI32ArrayAttr instData = layout.getInstData();
      tileShape = llvm::to_vector_of<int64_t>(instData.asArrayRef());