[Mlir-commits] [mlir] [MLIR][XeGPU] Add unroll patterns and blocking pass for XeGPU [2/N] (PR #140163)
Chao Chen
llvmlistbot at llvm.org
Tue May 27 08:06:13 PDT 2025
================
@@ -83,3 +100,275 @@ mlir::xegpu::getDistributedVectorType(VectorType originalType,
/*memory_space=*/xegpu::MemorySpace::Global, layout);
return xegpu::getDistributedVectorType(helperTdescTy);
}
+
+std::string xegpu::getLayoutName(OpOperand &opr) {
+ const StringRef prefix("layout_operand_");
+ return llvm::formatv("{0}{1}", prefix, opr.getOperandNumber()).str();
+}
+
+std::string xegpu::getLayoutName(OpResult res) {
+ const StringRef prefix = "layout_result_";
+ return llvm::formatv("{0}{1}", prefix, res.getResultNumber()).str();
+}
+
+xegpu::LayoutAttr xegpu::getLayoutAttr(Value value) {
+ if (!value)
+ return nullptr;
+
+ if (auto tdescTy =
+ dyn_cast_if_present<xegpu::TensorDescType>(value.getType()))
+ return tdescTy.getLayoutAttr();
+
+ if (auto result = dyn_cast<OpResult>(value)) {
+ Operation *defOp = result.getDefiningOp();
+ assert(defOp && "result must have a defining op");
+
+ // for LoadNdOp, the layout is stored in the tensor descriptor
+ if (auto loadNd = dyn_cast<xegpu::LoadNdOp>(defOp))
+ return getLayoutAttr(loadNd.getTensorDesc());
+
+ std::string layoutName = getLayoutName(result);
+ if (defOp->hasAttr(layoutName))
+ return defOp->getAttrOfType<xegpu::LayoutAttr>(layoutName);
+ }
+
+ if (auto arg = dyn_cast<BlockArgument>(value)) {
+ auto parentOp = arg.getOwner()->getParentOp();
+ if (auto loop = dyn_cast<LoopLikeOpInterface>(parentOp)) {
+ OpOperand *tiedInit = loop.getTiedLoopInit(arg);
+ return getLayoutAttr(tiedInit->get());
+ }
+ }
+
+ return nullptr;
+}
+
+xegpu::LayoutAttr xegpu::getLayoutAttr(OpOperand &opr) {
+ Operation *op = opr.getOwner();
+ std::string layoutName = xegpu::getLayoutName(opr);
+ if (op->hasAttr(layoutName))
+ return op->getAttrOfType<xegpu::LayoutAttr>(layoutName);
+ return getLayoutAttr(opr.get());
+}
+
+void xegpu::setLayoutAttr(OpOperand &opr, LayoutAttr layout) {
+ auto owner = opr.getOwner();
+ std::string name = xegpu::getLayoutName(opr);
+ if (layout && !owner->hasAttrOfType<LayoutAttr>(name))
+ owner->setAttr(name, layout);
+}
+
+void xegpu::setLayoutAttr(OpResult result, LayoutAttr layout) {
+ Operation *owner = result.getOwner();
+ std::string name = xegpu::getLayoutName(result);
+ if (layout && !owner->hasAttr(name))
+ owner->setAttr(name, layout);
+}
+
+void xegpu::setLayoutAttrs(Operation *mod,
+ function_ref<LayoutAttr(Value)> getLayoutImpl) {
+ mod->walk([&](Operation *op) {
+ for (OpResult result : op->getOpResults()) {
+ auto layout = getLayoutImpl(result);
+ setLayoutAttr(result, layout);
+ }
+ for (OpOperand &opr : op->getOpOperands()) {
+ auto layout = getLayoutImpl(opr.get());
+ setLayoutAttr(opr, layout);
+ }
+ });
+}
+
+SmallVector<Value>
+xegpu::extractVectorsWithShapeFromValue(OpBuilder &builder, Location loc,
+ Value value, ArrayRef<int64_t> shape) {
+ auto vecTy = dyn_cast<VectorType>(value.getType());
+ if (!vecTy)
+ return {value};
+
+ ArrayRef<int64_t> srcShape = vecTy.getShape();
+ if (!computeShapeRatio(srcShape, shape))
+ return {value};
+
+ SmallVector<Value> result;
+ for (SmallVector<int64_t> offsets : StaticTileOffsetRange(srcShape, shape)) {
+ SmallVector<int64_t> staticStrides(offsets.size(), 1);
+ result.push_back(builder.create<vector::ExtractStridedSliceOp>(
+ loc, value, offsets, shape, staticStrides));
+ }
+
+ return result;
+}
+
+Value xegpu::createVectorWithShapeFromValues(OpBuilder &builder, Location loc,
+ ValueRange values,
+ ArrayRef<int64_t> shape) {
+ VectorType inputTy = dyn_cast<VectorType>(values[0].getType());
+ assert(llvm::all_of(values.getTypes(),
+ [&](Type type) { return type == inputTy; }) &&
+ "values must be of the same VectorType");
+
+ Type elemTy = inputTy.getElementType();
+ ArrayRef<int64_t> tileShape = inputTy.getShape();
+
+ VectorType resultTy = VectorType::get(shape, elemTy);
+ auto zeroAttr = builder.getZeroAttr(elemTy);
+ Value result = builder.create<arith::ConstantOp>(
+ loc, resultTy, DenseElementsAttr::get(resultTy, zeroAttr));
+
+ for (auto [src, offsets] :
+ llvm::zip_equal(values, StaticTileOffsetRange(shape, tileShape))) {
+ SmallVector<int64_t> staticStrides(offsets.size(), 1);
+ result = builder.create<vector::InsertStridedSliceOp>(
+ loc, src, result, offsets, staticStrides);
+ }
+ return result;
+}
+
+void xegpu::doSCFStructuralTypeConversionWithTensorType(
+ Operation *op, TypeConverter converter) {
+ MLIRContext *context = op->getContext();
+
+ auto materializeCast = [&](OpBuilder &builder, Type type, ValueRange inputs,
+ Location loc) -> Value {
+ return builder.create<UnrealizedConversionCastOp>(loc, type, inputs)
+ .getResult(0);
+ };
+
+ { // convert VectorType to RankedTensorType for SCF Structural ops
+ TypeConverter converter;
+ converter.addConversion([&](Type type) -> Type { return type; });
+ converter.addConversion([&](VectorType type) -> Type {
+ return RankedTensorType::get(type.getShape(), type.getElementType());
+ });
+ converter.addSourceMaterialization(materializeCast);
+ converter.addTargetMaterialization(materializeCast);
+
+ mlir::ConversionTarget target(*context);
+ target.addLegalOp<UnrealizedConversionCastOp>();
+
+ mlir::RewritePatternSet patterns(context);
+ scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
+ target);
+ (void)mlir::applyPartialConversion(op, target, std::move(patterns));
+ }
+
+ { // propagate the layout attribute to RankedTensorType by checking
+ // BuiltInUnrealizedCastOps
+ // for VectorType to RankedTensorType cast.
+ op->walk([&](UnrealizedConversionCastOp castOp) {
+ if (castOp.getNumOperands() != 1 || castOp.getNumResults() != 1)
+ return WalkResult::skip();
+
+ Value input = castOp.getInputs()[0];
+ Value result = castOp.getResults()[0];
+ auto inputTy = dyn_cast<VectorType>(input.getType());
+ auto resultTy = dyn_cast<RankedTensorType>(result.getType());
+
+ // Only look at ops casting from VectorType to RankedTensorType
+ if (!isa<VectorType>(inputTy) || !isa<RankedTensorType>(resultTy))
+ return WalkResult::skip();
+
+ xegpu::LayoutAttr layout = xegpu::getLayoutAttr(input);
+ if (!layout)
+ return WalkResult::skip();
+
+ RankedTensorType newTy = resultTy.cloneWithEncoding(layout);
+ result.setType(newTy);
+
+ // update the arguments if user is a LoopLike op.
+ for (OpOperand &use : result.getUses()) {
+ if (auto loop = dyn_cast<LoopLikeOpInterface>(use.getOwner())) {
+ BlockArgument arg = loop.getTiedLoopRegionIterArg(&use);
+ arg.setType(newTy);
+ }
+ // whileOp has two regions, the BlockArgument of the after region
+ // is not exposed by LoopLikeOpInterface
+ if (auto whileOp = dyn_cast<scf::WhileOp>(use.getOwner())) {
+ unsigned idx = use.getOperandNumber();
+ BlockArgument arg = whileOp.getAfterArguments()[idx];
+ arg.setType(newTy);
+ }
+ }
+ return WalkResult::advance();
+ });
+
+ // using yieldOp as anchor to update the result type of its ParentOp
+ op->walk([&](scf::YieldOp yieldOp) {
+ Operation *parentOp = yieldOp->getParentOp();
+ for (OpResult r : parentOp->getOpResults()) {
+ unsigned idx = r.getResultNumber();
+ Type resultTy = r.getType();
+ Type yieldTy = yieldOp.getResults()[idx].getType();
+ if (isa<RankedTensorType>(resultTy) && yieldTy != resultTy)
+ r.setType(yieldTy);
+ }
+ });
+ }
+
+ { // perform the conversion from RankedTensorType to VectorType based on the
+ // LayoutAttr
+
+ // Handle the UnrealizedConversionCastOp introduced by the first step.
+ // For vector->RankedTensorType, it will simply forward the inputs.
+ // For RankedTensorType->vector, it will update the inputs with the
+ // one from the adaptor.
+ class UnrealizedConversionCastOpPattern
+ : public OpConversionPattern<mlir::UnrealizedConversionCastOp> {
+ using OpConversionPattern<
+ mlir::UnrealizedConversionCastOp>::OpConversionPattern;
+
+ mlir::LogicalResult
+ matchAndRewrite(mlir::UnrealizedConversionCastOp op,
+ OneToNOpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto inputs = op.getOperands();
+ auto outputs = op.getOutputs();
+
+ if (inputs.size() != 1 || outputs.size() != 1)
+ return failure();
+
+ auto inputTy = inputs[0].getType();
+ auto outputTy = outputs[0].getType();
+
+ if (isa<VectorType>(inputTy) && isa<RankedTensorType>(outputTy)) {
+ rewriter.replaceOpWithMultiple(op, adaptor.getInputs());
+ return success();
+ }
+
+ if (isa<RankedTensorType>(inputTy) && isa<VectorType>(outputTy)) {
+ SmallVector<Value> values = flattenValues(adaptor.getInputs());
+ auto newOp = rewriter.create<UnrealizedConversionCastOp>(
+ op.getLoc(), outputTy, values);
+ rewriter.replaceOp(op, newOp);
+ return success();
+ }
+ return failure();
+ }
+ };
+
+ converter.addSourceMaterialization(materializeCast);
+ converter.addTargetMaterialization([&](OpBuilder &builder, TypeRange type,
----------------
chencha3 wrote:
I tried, it doesn't work. Majorly because here I expect the `TypeRange` instead of `Type`?
https://github.com/llvm/llvm-project/pull/140163
More information about the Mlir-commits
mailing list