Correctly emit subindex for pointer offset

Author: ivanradanov

tools/cgeist/Lib/clang-mlir.cc

@@ -2896,6 +2896,18 @@ ValueCategory MLIRScanner::VisitBinaryOperator(clang::BinaryOperator *BO) {
     }
   }
   case clang::BinaryOperator::Opcode::BO_Add: {
+    auto emitSubindex = [&](auto mr, auto ptradd) {
+      auto mt = mr.getType().template dyn_cast<mlir::MemRefType>();
+      auto shape = std::vector<int64_t>(mt.getShape());
+      shape[0] = -1;
+      auto mt0 = mlir::MemRefType::get(shape, mt.getElementType(),
+                                       MemRefLayoutAttrInterface(),
+                                       mt.getMemorySpace());
+      ptradd = castToIndex(loc, ptradd);
+      return ValueCategory(
+          builder.create<polygeist::SubIndexOp>(loc, mt0, mr, ptradd),
+          /*isReference*/ false);
+    };
     if (auto cty = dyn_cast<clang::ComplexType>(BO->getType())) {
       mlir::Value real =
           builder.create<AddFOp>(loc, getComplexPart(loc, lhs.val, 0),
@@ -2911,17 +2923,10 @@ ValueCategory MLIRScanner::VisitBinaryOperator(clang::BinaryOperator *BO) {
     if (lhs_v.getType().isa<mlir::FloatType>()) {
       return ValueCategory(builder.create<AddFOp>(loc, lhs_v, rhs_v),
                            /*isReference*/ false);
-    } else if (auto mt = lhs_v.getType().dyn_cast<mlir::MemRefType>()) {
-      auto shape = std::vector<int64_t>(mt.getShape());
-      shape[0] = -1;
-      auto mt0 = mlir::MemRefType::get(shape, mt.getElementType(),
-                                       MemRefLayoutAttrInterface(),
-                                       mt.getMemorySpace());
-      auto ptradd = rhs_v;
-      ptradd = castToIndex(loc, ptradd);
-      return ValueCategory(
-          builder.create<polygeist::SubIndexOp>(loc, mt0, lhs_v, ptradd),
-          /*isReference*/ false);
+    } else if (lhs_v.getType().isa<mlir::MemRefType>()) {
+      return emitSubindex(lhs_v, rhs_v);
+    } else if (rhs_v.getType().isa<mlir::MemRefType>()) {
+      return emitSubindex(rhs_v, lhs_v);
     } else if (auto pt =
                    lhs_v.getType().dyn_cast<mlir::LLVM::LLVMPointerType>()) {
       return ValueCategory(

tools/cgeist/Test/Verification/memrefaddassign.cpp

@@ -4,9 +4,26 @@ float *foo(float *a) {
 	a += 32;
 	return a;
 }
+float *foo1(float *a) {
+	return a + 32;
+}
+float *foo2(float *a) {
+	return 32 + a;
+}
 // CHECK: func @_Z3fooPf(%[[arg0:.+]]: memref<?xf32>)
 // CHECK-NEXT   %[[c32:.+]] = arith.constant 32 : index
 // CHECK-NEXT   %[[V0:.+]] = "polygeist.subindex"(%[[arg0]], %[[c32]]) : (memref<?xf32>, index) -> memref<?xf32>
 // CHECK-NEXT   return %[[V0]] : memref<?xf32>
 // CHECK-NEXT }
 
+// CHECK: func @_Z4foo1Pf(%[[arg0:.+]]: memref<?xf32>)
+// CHECK-NEXT   %[[c32:.+]] = arith.constant 32 : index
+// CHECK-NEXT   %[[V0:.+]] = "polygeist.subindex"(%[[arg0]], %[[c32]]) : (memref<?xf32>, index) -> memref<?xf32>
+// CHECK-NEXT   return %[[V0]] : memref<?xf32>
+// CHECK-NEXT }
+
+// CHECK: func @_Z4foo2Pf(%[[arg0:.+]]: memref<?xf32>)
+// CHECK-NEXT   %[[c32:.+]] = arith.constant 32 : index
+// CHECK-NEXT   %[[V0:.+]] = "polygeist.subindex"(%[[arg0]], %[[c32]]) : (memref<?xf32>, index) -> memref<?xf32>
+// CHECK-NEXT   return %[[V0]] : memref<?xf32>
+// CHECK-NEXT }