Implemented improvement in linalg debufferize to work through inversesubmap, and a more shophisticated approach taken for tracking current tensors via a tree of regions

Author: arpitj1

include/polygeist/PolygeistOps.td

@@ -260,19 +260,111 @@ def TypeAlignOp : Polygeist_Op<"typeAlign", [Pure]> {
   let hasCanonicalizer = 1;
 }
 
+//Add check for result to be same as original memref/tensor type
+def SubmapInverseOp : Polygeist_Op<"submapInverse", [Pure, ViewLikeOpInterface]> {
+  let summary = "Inverse submap operation for scatter-back semantics";
+  let description = [{
+    The `polygeist.submapInverse` operation scatters a modified view back into
+    the original base tensor/memref, preserving elements not covered by the view.
+    
+    This is the inverse operation to `polygeist.submap` and is essential for
+    debufferization of strided memory operations.
+    
+    Example:
+    ```mlir
+    // Scatter strided view back into base tensor
+    %base_updated = polygeist.submapInverse(%base, %modified_view, %stride, %size)
+      <{map = affine_map<(d0)[s0] -> (d0 * s0)>}>
+      : (tensor<100xf32>, tensor<50xf32>) -> tensor<100xf32>
+    
+    // Semantics: base_updated[i*stride] = modified_view[i]
+    //            base_updated[other] = base[other] (preserved)
+    ```
+  }];
+  
+  let arguments = (ins 
+      Arg<AnyTypeOf<[AnyMemRef, AnyTensor]>, "the original base">:$base_original,
+      Arg<AnyTypeOf<[AnyMemRef, AnyTensor]>, "the modified view">:$view_modified,
+      Variadic<Index>:$indices_and_sizes,
+      AffineMapAttr:$map
+      );
+  let results = (outs AnyTypeOf<[AnyMemRef, AnyTensor]> : $result);
+  let hasFolder = 1;
+  let hasCanonicalizer = 1;
+  
+  let assemblyFormat = [{
+    `(` $base_original `,` $view_modified (`,` $indices_and_sizes^)? `)` 
+    attr-dict `:` functional-type(operands, results)
+  }];
+
+  let extraClassDeclaration = [{
+    ::mlir::ValueRange getSymbols() { return getOperands().slice(2, getMap().getNumSymbols()); }
+    ::mlir::ValueRange getSizes() { 
+      auto shapedType = ::llvm::cast<::mlir::ShapedType>(getType());
+      return getOperands().slice(getMap().getNumSymbols()+2, shapedType.getShape().size()); 
+    }
+    ::mlir::Value getViewSource() { return getBaseOriginal(); }
+    
+    // Type compatibility helpers
+    bool isMemRefVariant() { 
+      return ::llvm::isa<::mlir::MemRefType>(getBaseOriginal().getType());
+    }
+    bool isTensorVariant() { 
+      return ::llvm::isa<::mlir::TensorType>(getBaseOriginal().getType());
+    }
+  }];
+}
+
 def SubmapOp : Polygeist_Op<"submap", [Pure, ViewLikeOpInterface]> {
-  let arguments = (ins Arg<AnyMemRef, "the reference to load from">:$memref,
+  let summary = "Submap operation for strided view extraction";
+  let description = [{
+    The `polygeist.submap` operation creates a strided view of a tensor/memref
+    by applying an affine map to extract elements. This is used to represent
+    strided access patterns in a composable way.
+    
+    The operation works in both memref and tensor contexts, enabling
+    debufferization of strided operations.
+    
+    Example:
+    ```mlir
+    // Extract every other element (stride=2)
+    %view = polygeist.submap(%base, %stride, %size) 
+            <{map = affine_map<(d0)[s0] -> (d0 * s0)>}>
+            : tensor<100xf32> -> tensor<50xf32>
+    
+    // Semantics: view[i] = base[i * stride]
+    ```
+  }];
+  
+  let arguments = (ins 
+      Arg<AnyTypeOf<[AnyMemRef, AnyTensor]>, "the base to view">:$base,
       Variadic<Index>:$indices_and_sizes,
       AffineMapAttr:$map
       );
-  let results = (outs AnyMemRef : $result);
+  let results = (outs AnyTypeOf<[AnyMemRef, AnyTensor]> : $result);
   let hasFolder = 1;
   let hasCanonicalizer = 1;
+  
+  let assemblyFormat = [{
+    `(` $base (`,` $indices_and_sizes^)? `)` 
+    attr-dict `:` functional-type(operands, results)
+  }];
 
   let extraClassDeclaration = [{
     ::mlir::ValueRange getSymbols() { return getOperands().slice(1, getMap().getNumSymbols()); }
-    ::mlir::ValueRange getSizes() { return getOperands().slice(getMap().getNumSymbols()+1, getType().getShape().size()); }
-    ::mlir::Value getViewSource() { return getMemref(); }
+    ::mlir::ValueRange getSizes() { 
+      auto shapedType = ::llvm::cast<::mlir::ShapedType>(getType());
+      return getOperands().slice(getMap().getNumSymbols()+1, shapedType.getShape().size()); 
+    }
+    ::mlir::Value getViewSource() { return getBase(); }
+    
+    // Type compatibility helpers
+    bool isMemRefVariant() { 
+      return ::llvm::isa<::mlir::MemRefType>(getBase().getType());
+    }
+    bool isTensorVariant() { 
+      return ::llvm::isa<::mlir::TensorType>(getBase().getType());
+    }
   }];
 }
 

lib/polygeist/Ops.cpp

@@ -5769,7 +5769,7 @@ struct SubMapOpCanonicalize : public OpRewritePattern<polygeist::SubmapOp> {
     /// %x = ... : memref<4x5xf32>
     //  %y = memref.cast %x : memref<4x5xf32> -> memref<?x?xf32>
     //
-    auto source_memref = op.getMemref();
+    auto source_memref = op.getBase();
     bool isIdentity = op.getMap().isIdentity();
     bool isInputSameDim = llvm::all_of(
         llvm::zip_equal(op.getSizes(),
@@ -5785,7 +5785,7 @@ struct SubMapOpCanonicalize : public OpRewritePattern<polygeist::SubmapOp> {
         });
     if (isIdentity && isInputSameDim) {
       rewriter.replaceOpWithNewOp<memref::CastOp>(op, op.getType(),
-                                                  op.getMemref());
+                                                  op.getBase());
       return success();
     }
     if (auto sapOp = source_memref.getDefiningOp<polygeist::SubmapOp>()) {
@@ -5797,7 +5797,7 @@ struct SubMapOpCanonicalize : public OpRewritePattern<polygeist::SubmapOp> {
       operands.append(op.getSymbols().begin(), op.getSymbols().end());
       operands.append(op.getSizes().begin(), op.getSizes().end());
       rewriter.replaceOpWithNewOp<polygeist::SubmapOp>(
-          op, op.getType(), sapOp.getMemref(), operands, new_map);
+          op, op.getType(), sapOp.getBase(), operands, new_map);
       return success();
     }
     return failure();
@@ -5990,7 +5990,7 @@ static bool canConvertSubmapToSubView(polygeist::SubmapOp submapOp) {
   auto map = submapOp.getMap();
   auto sizes = submapOp.getSizes();
   auto symbols = submapOp.getSymbols();
-  auto source_memref = submapOp.getMemref();
+  auto source_memref = submapOp.getBase();
 
   // 0. Only convert if map has symbols
   if (submapOp.getMap().getNumSymbols() == 0) {
@@ -6111,7 +6111,7 @@ struct SubmapToSubviewOp : public OpRewritePattern<polygeist::SubmapOp> {
     for (Value size : conversionInfo.sizes) {
       sizeValues.push_back(size);
     }
-    rewriter.replaceOpWithNewOp<memref::SubViewOp>(submapOp, submapOp.getType(), submapOp.getMemref(), offsetValues, sizeValues, strideValues);
+    rewriter.replaceOpWithNewOp<memref::SubViewOp>(submapOp, submapOp.getBase(), offsetValues, sizeValues, strideValues);
     return success();
   }
 };
@@ -6535,7 +6535,7 @@ class LoadSubMap final : public OpRewritePattern<affine::AffineLoadOp> {
 
     auto submap_map = subMapOp.getMap();
     auto submap_operands = subMapOp.getSymbols();
-    auto source_memref = subMapOp.getMemref();
+    auto source_memref = subMapOp.getBase();
 
     auto load_map = op.getAffineMap();
     auto load_operands = op.getMapOperands();
@@ -6567,7 +6567,7 @@ class StoreSubMap final : public OpRewritePattern<affine::AffineStoreOp> {
 
     auto submap_map = subMapOp.getMap();
     auto submap_operands = subMapOp.getSymbols();
-    auto source_memref = subMapOp.getMemref();
+    auto source_memref = subMapOp.getBase();
 
     auto load_map = op.getAffineMap();
     auto load_operands = op.getMapOperands();
@@ -6705,3 +6705,20 @@ void polygeist::SubmapOp::getCanonicalizationPatterns(
   // results.insert<LoadSubMap, StoreSubMap, DimSubMap>(context);
 }
 
+//===----------------------------------------------------------------------===//
+// SubmapInverseOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult mlir::polygeist::SubmapInverseOp::fold(
+    mlir::polygeist::SubmapInverseOp::FoldAdaptor adaptor) {
+  // TODO: Add folding logic for SubmapInverseOp
+  // For now, just return nullptr (no folding)
+  return nullptr;
+}
+
+void polygeist::SubmapInverseOp::getCanonicalizationPatterns(
+    RewritePatternSet &results, MLIRContext *context) {
+  // TODO: Add canonicalization patterns for SubmapInverseOp
+  // For now, leave empty
+}
+