[SYCL-MLIR] Fix memref2pointer(subindex) optimization (#7408)

Optimization `Memref2PointerIndex` is only valid when the element type
of the source of subindex is not a struct type.
After this PR, `stream-copy.cpp` and `stream-triad.cpp` can both run
successfully. They are added in
intel/llvm-test-suite#1384.

Signed-off-by: Tsang, Whitney <[email protected]>

Author: whitneywhtsang

polygeist/lib/polygeist/Ops.cpp

@@ -1159,7 +1159,7 @@ class Memref2Pointer2MemrefCast final
     return success();
   }
 };
-/// Simplify pointer2memref(memref2pointer(x)) to cast(x)
+/// Simplify memref2pointer(subindex(x)) to getelementptr(memref2pointer(x))
 class Memref2PointerIndex final : public OpRewritePattern<Memref2PointerOp> {
 public:
   using OpRewritePattern<Memref2PointerOp>::OpRewritePattern;
@@ -1173,9 +1173,12 @@ class Memref2PointerIndex final : public OpRewritePattern<Memref2PointerOp> {
     if (src.getSource().getType().cast<MemRefType>().getShape().size() != 1)
       return failure();
 
+    auto MET = src.getSource().getType().cast<MemRefType>().getElementType();
+    if (MET.isa<LLVM::LLVMStructType>())
+      return failure();
+
     Value idx[] = {src.getIndex()};
     auto PET = op.getType().cast<LLVM::LLVMPointerType>().getElementType();
-    auto MET = src.getSource().getType().cast<MemRefType>().getElementType();
     if (PET != MET) {
       auto ps = rewriter.create<polygeist::TypeSizeOp>(
           op.getLoc(), rewriter.getIndexType(), mlir::TypeAttr::get(PET));

polygeist/test/polygeist-opt/invalid_canonicalization.mlir

@@ -35,7 +35,7 @@ func.func @SubToCast(%arg0: memref<?x!llvm.struct<(i32)>>) -> memref<?xi32> {
 // CHECK-NEXT:     [[C1:%.*]] = arith.constant 0 : i32
 // CHECK-NEXT:     [[T0:%.*]] = "polygeist.subindex"([[A0]], [[C0]]) : (memref<?x!llvm.struct<(i32)>>, index) -> memref<?xi32>
 // CHECK-NEXT:     memref.store [[C1]], [[T0]][[[C0]]] : memref<?xi32>
-// CHECK-NEXT:     return %0 : memref<?xi32>
+// CHECK-NEXT:     return [[T0]] : memref<?xi32>
 // CHECK-NEXT:  }
 func.func @SimplifySubIndexUsers(%arg0: memref<?x!llvm.struct<(i32)>>) -> memref<?xi32> {
   %c0 = arith.constant 0 : index
@@ -46,3 +46,19 @@ func.func @SimplifySubIndexUsers(%arg0: memref<?x!llvm.struct<(i32)>>) -> memref
 }
 
 // -----
+
+// CHECK:  func.func @Memref2PointerIndex([[A0:%.*]]: memref<?x!llvm.struct<(i32, i32)>>) -> !llvm.ptr<i32> {
+// CHECK-NEXT:    [[C1:%.*]] = arith.constant 1 : index
+// CHECK-NEXT:    [[T0:%.*]] = "polygeist.subindex"([[A0]], [[C1]]) : (memref<?x!llvm.struct<(i32, i32)>>, index) -> memref<?xi32>
+// CHECK-NEXT:    [[T1:%.*]] = "polygeist.memref2pointer"([[T0]]) : (memref<?xi32>) -> !llvm.ptr<i32>
+// CHECK-NEXT:    return [[T1]] : !llvm.ptr<i32>
+// CHECK-NEXT:  }
+
+func.func @Memref2PointerIndex(%arg0: memref<?x!llvm.struct<(i32, i32)>>) -> !llvm.ptr<i32> {
+  %c1 = arith.constant 1 : index
+  %0 = "polygeist.subindex"(%arg0, %c1) : (memref<?x!llvm.struct<(i32, i32)>>, index) -> memref<?xi32>
+  %1 = "polygeist.memref2pointer"(%0) : (memref<?xi32>) -> !llvm.ptr<i32>
+  return %1 : !llvm.ptr<i32>
+}
+
+// -----