Add support for non memref source for init tile in wg to sg pass (#1060)

Add support for non memref source for init tile

Author: nbpatel

lib/Dialect/XeTile/Transforms/WgToSg.cpp

@@ -208,20 +208,18 @@ class WGToSGInitTileOpPattern : public OpConversionPattern<xetile::InitTileOp> {
     for (size_t j = 0; j < offsets.size() - 2; ++j) {
         newOffsets.push_back(offsets[j]);
     }
+
+    auto sourceMemRefType = mlir::dyn_cast<mlir::MemRefType>(source.getType());
+
     for (size_t i = 0; i < offsetPermutations.size(); i++) {
       newOffsets.push_back(offsetPermutations[i][0]);
       newOffsets.push_back(offsetPermutations[i][1]);
       Value newInitTileOp = nullptr;
-      auto sourceMemRefType = mlir::dyn_cast<mlir::MemRefType>(source.getType());
-      if (!sourceMemRefType) {
-        return failure();
-      }
-
-      if (sourceMemRefType.hasStaticShape()) {
+      if (sourceMemRefType && sourceMemRefType.hasStaticShape()) {
         newInitTileOp = rewriter.create<xetile::InitTileOp>(
             loc, newTileTy, source, newOffsets);
       }
-      else {
+      else { // memref with dynamic shape or non memref source
         newInitTileOp = rewriter.create<xetile::InitTileOp>(
             loc, newTileTy, source, newOffsets, op.getMixedSizes(), op.getMixedStrides());
       }

test/Dialect/XeTile/Transforms/WgToSg/unit_tests.mlir

@@ -160,4 +160,12 @@ gpu.module @test_arith_extf {
         }
         gpu.return
     }
+
+    gpu.func @test_init_tile_using_addr(%src: i64, %dim0_size : index, %dim1_size : index,
+      %dim0_stride : index, %dim1_stride : index ) {
+      //CHECK: xetile.init_tile {{%.*}}[{{%.*}}, {{%.*}}], [{{%.*}}, {{%.*}}], [{{%.*}}, {{%.*}}] : i64 -> !xetile.tile<4x64xf16>
+      %1 = xetile.init_tile %src[8, 16], [%dim0_size, %dim1_size], [%dim0_stride, %dim1_stride]
+        : i64 ->  !xetile.tile<128x64xf16, #xetile.tile_attr<wg_map = <sg_layout = [32, 1], sg_data = [4, 64]>>>
+      gpu.return
+  }
 }