partial working

Author: chencha3

mlir/lib/Dialect/XeGPU/Transforms/XeGPUInstructionlize.cpp

@@ -120,18 +120,22 @@ bool XeGPUInstructionlizePass::needsUnroll(Operation *op) const {
 
 void XeGPUInstructionlizePass::runOnOperation() {
   MLIRContext *ctx = &getContext();
+  Operation *op = getOperation();
+
+  // first perform type conversion for SCF control folow ops
+  xegpu::doSCFStructuralTypeConversionWithTensorType(op);
+
   xegpu::UnrollOptions options;
   options.setFilterConstraint([&](Operation *op) -> LogicalResult {
     return needsUnroll(op) ? success() : failure();
   });
 
-  options.setNativeShapeFn(
-      [&](Operation *op) -> std::optional<SmallVector<int64_t>> {
+  options.setNativeShapeFn([&](Operation *op) {
         return getTileShape(op);
       });
 
   options.setUnrolledTypesFn(
-      [&](ShapedType type, ArrayRef<int64_t> tileShape) -> SmallVector<Type> {
+      [&](ShapedType type, ArrayRef<int64_t> tileShape) {
         Type elemTy = type.getElementType();
         Type newTy;
 
@@ -149,8 +153,10 @@ void XeGPUInstructionlizePass::runOnOperation() {
         return SmallVector<Type>(computeProduct(*ratio), newTy);
       });
 
-  RewritePatternSet patterns(ctx);
+  GreedyRewriteConfig config;
+  config.setStrictness(GreedyRewriteStrictness::ExistingOps);
 
+  RewritePatternSet patterns(ctx);
   populateXeGPUUnrollPatterns(patterns, options);
-  (void)applyPatternsGreedily(getOperation(), std::move(patterns));
+  (void)applyPatternsGreedily(getOperation(), std::move(patterns), config);
 }

mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp

@@ -215,17 +215,16 @@ void xegpu::doSCFStructuralTypeConversionWithTensorType(Operation *op) {
     // LayoutAttr
 
     auto computeTileShapeAndCount = [&](ArrayRef<int64_t> shape,
-                                        DenseI32ArrayAttr sgDataAttr,
-                                        DenseI32ArrayAttr sgLayoutAttr) {
+                                          DenseI32ArrayAttr sgDataAttr,
+                                          DenseI32ArrayAttr sgLayoutAttr) {
       SmallVector<int64_t> tileShape;
       auto sgLayout = llvm::to_vector_of<int64_t>(sgLayoutAttr.asArrayRef());
       if (sgDataAttr)
         tileShape = llvm::to_vector_of<int64_t>(sgDataAttr.asArrayRef());
       else
         tileShape = computeShapeRatio(shape, sgLayout).value_or(tileShape);
       assert(tileShape.size() && "failed to compute tileShape");
-      SmallVector<int64_t> distUnit =
-          computeElementwiseMul(sgLayout, tileShape);
+      SmallVector<int64_t> distUnit = computeElementwiseMul(sgLayout, tileShape);
       int count = computeProduct(shape) / computeProduct(distUnit);
       return std::make_pair(tileShape, count);
     };
@@ -249,8 +248,7 @@ void xegpu::doSCFStructuralTypeConversionWithTensorType(Operation *op) {
             if (layout.isWgLayout()) {
               // for WgToSg, the subShape is either from sgData or computed as
               // shape/sgLayout
-              std::tie(subShape, count) = computeTileShapeAndCount(
-                  shape, layout.getSgData(), layout.getSgLayout());
+              std::tie(subShape, count) = computeTileShapeAndCount(shape, layout.getSgData(), layout.getSgLayout());
             } else if (DenseI32ArrayAttr instData = layout.getInstData()) {
               // for unrolling, the subShape is determined by inst_data
               subShape = llvm::to_vector_of<int64_t>(instData.asArrayRef());
@@ -280,8 +278,7 @@ void xegpu::doSCFStructuralTypeConversionWithTensorType(Operation *op) {
             if (layout.isWgLayout()) {
               // for WgToSg, the subShape is either from sgData or computed as
               // shape/sgLayout
-              std::tie(subShape, count) = computeTileShapeAndCount(
-                  shape, layout.getSgData(), layout.getSgLayout());
+              std::tie(subShape, count) = computeTileShapeAndCount(shape, layout.getSgData(), layout.getSgLayout());
               layout = layout.dropSgLayoutAndData();
             } else if (DenseI32ArrayAttr instData = layout.getInstData()) {
               // for unrolling, the subShape is determined by inst_data
@@ -298,7 +295,11 @@ void xegpu::doSCFStructuralTypeConversionWithTensorType(Operation *op) {
         });
 
     converter.addSourceMaterialization(materializeCast);
-    converter.addTargetMaterialization(materializeCast);
+    converter.addTargetMaterialization([&](OpBuilder &builder, TypeRange type,
+                                        ValueRange inputs, Location loc) {
+      return builder.create<UnrealizedConversionCastOp>(loc, type, inputs)
+          .getResults();
+    });
 
     mlir::ConversionTarget target(*context);
     target.addLegalOp<UnrealizedConversionCastOp>();