[NFC] Rename load_from/store_to_memref to load_from/store_to_buffer (#20897)

Renames `iree_codegen.load_from_memref` to
`iree_codegen.load_from_buffer`, and `iree_codegen.store_to_memref` to
`iree_codegen.store_to_buffer` for consistency with the op definition.

Signed-off-by: Max Dawkins <[email protected]>

Author: Max191

compiler/src/iree/compiler/Codegen/Common/BufferizeDispatchTensorLoadStore.cpp

@@ -56,7 +56,7 @@ bufferizeDispatchTensorLoad(RewriterBase &rewriter,
       rewriter, loadOp.getLoc(), subspanOp, loadOp.getType(),
       loadOp.getMixedOffsets(), loadOp.getMixedSizes(),
       loadOp.getMixedStrides());
-  rewriter.replaceOpWithNewOp<IREE::Codegen::LoadFromMemrefOp>(
+  rewriter.replaceOpWithNewOp<IREE::Codegen::LoadFromBufferOp>(
       loadOp, loadOp.getType(), sourceBuffer);
 }
 
@@ -68,10 +68,10 @@ bufferizeDispatchTensorStore(RewriterBase &rewriter,
   if (!subspanOp) {
     return;
   }
-  // For the store_to_memref op, generate any subviews as early as possible in
-  // the IR. This opens more opportunities for using the store_to_memref op's
+  // For the store_to_buffer op, generate any subviews as early as possible in
+  // the IR. This opens more opportunities for using the store_to_buffer op's
   // SubsetInsertionOpInterface, since equivalent subset extractions can only be
-  // created after the store_to_memref op's output (the subspan or subview) in
+  // created after the store_to_buffer op's output (the subspan or subview) in
   // the IR.
   OpBuilder::InsertionGuard g(rewriter);
   (void)setInsertionPointAfterLastNeededValue(
@@ -83,7 +83,7 @@ bufferizeDispatchTensorStore(RewriterBase &rewriter,
 
   Value tensor = storeOp.getValue();
   rewriter.setInsertionPoint(storeOp);
-  rewriter.replaceOpWithNewOp<IREE::Codegen::StoreToMemrefOp>(storeOp, tensor,
+  rewriter.replaceOpWithNewOp<IREE::Codegen::StoreToBufferOp>(storeOp, tensor,
                                                               outputBuffer);
 }
 

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp

@@ -259,7 +259,7 @@ static void buildNestedDistributionLoops(
 /// the writing of padding values into a separate operation on the buffer that
 /// the map_scatter op is ultimately written into. The result buffer is taken
 /// from the direct consumer of the `mapScatterOp`, which is expected to be an
-/// `iree_codegen.store_to_memref` op. Return failure if the result buffer is
+/// `iree_codegen.store_to_buffer` op. Return failure if the result buffer is
 /// not found.
 static FailureOr<MapScatterOp>
 foldPadIntoMapScatter(RewriterBase &rewriter, tensor::PadOp padOp,
@@ -270,12 +270,12 @@ foldPadIntoMapScatter(RewriterBase &rewriter, tensor::PadOp padOp,
     return rewriter.notifyMatchFailure(
         mapScatterOp, "map_scatter does not have a single user");
   }
-  auto storeOp = dyn_cast<IREE::Codegen::StoreToMemrefOp>(
+  auto storeOp = dyn_cast<IREE::Codegen::StoreToBufferOp>(
       *mapScatterOp->getUsers().begin());
   if (!storeOp) {
     return rewriter.notifyMatchFailure(
         mapScatterOp,
-        "map_scatter user is not an iree_codegen.store_to_memref op");
+        "map_scatter user is not an iree_codegen.store_to_buffer op");
   }
 
   rewriter.setInsertionPointAfter(storeOp);
@@ -420,7 +420,7 @@ combineRelayoutOpChain(RewriterBase &rewriter, MapScatterOp mapScatterOp,
 
 static MapScatterOp
 insertIdentityMapScatter(RewriterBase &rewriter,
-                         IREE::Codegen::StoreToMemrefOp storeOp) {
+                         IREE::Codegen::StoreToBufferOp storeOp) {
   Location loc = storeOp->getLoc();
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(storeOp);
@@ -447,11 +447,11 @@ combineLayoutTransformation(MLIRContext *ctx, FunctionOpInterface funcOp,
   IRRewriter rewriter(ctx);
   simplifyComplexRelayoutOps(rewriter, funcOp);
 
-  // Start from iree_codegen.store_to_memref ops, and combine producer
+  // Start from iree_codegen.store_to_buffer ops, and combine producer
   // relayout ops into a single map_scatter.
-  SmallVector<IREE::Codegen::StoreToMemrefOp> dispatchResults(
-      funcOp.getFunctionBody().getOps<IREE::Codegen::StoreToMemrefOp>());
-  for (IREE::Codegen::StoreToMemrefOp dispatchResult : dispatchResults) {
+  SmallVector<IREE::Codegen::StoreToBufferOp> dispatchResults(
+      funcOp.getFunctionBody().getOps<IREE::Codegen::StoreToBufferOp>());
+  for (IREE::Codegen::StoreToBufferOp dispatchResult : dispatchResults) {
     MapScatterOp mapScatterOp =
         insertIdentityMapScatter(rewriter, dispatchResult);
     combineRelayoutOpChain(rewriter, mapScatterOp, padDistributionConfigFn);

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.h

@@ -32,15 +32,15 @@ using PadDistributionConfigFn = function_ref<SmallVector<DistributionConfig>(
     ArrayRef<int64_t> iterationBounds, MLIRContext *)>;
 
 /// Combines any layout/indexing transformation ops at the ends of a dispatch.
-/// Finds `iree_codegen.store_to_memref` ops in the `funcOp`, and combines any
+/// Finds `iree_codegen.store_to_buffer` ops in the `funcOp`, and combines any
 /// layout transformation ops (like expand_shape, transpose, pack, etc.) that
 /// produce the tensor being stored into a single `iree_linalg_ext.map_scatter`
 /// op.
 ///
 /// This transformation will also combine `tensor.pad` ops into the map_scatter
-/// op, by moving the writing of the padding values to after the store_to_memref
+/// op, by moving the writing of the padding values to after the store_to_buffer
 /// op, and writing the padding values directly to the output buffer of the
-/// store_to_memref. The writes of the pad values will be distributed based on
+/// store_to_buffer. The writes of the pad values will be distributed based on
 /// the `DistributionConfig`s returned by `padDistributionConfigFn`, and then
 /// the inner distributed tile will be tiled to a loop nest of memref.store ops.
 LogicalResult

compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td

@@ -35,12 +35,12 @@ def GPUCombineLayoutTransformationPass :
   let summary =
     "Combines layout transformation operations into a single map_scatter operation.";
   let description = [{
-    Starting from iree_codegen.store_to_memref ops, iteratively combine producer
+    Starting from iree_codegen.store_to_buffer ops, iteratively combine producer
     layout/indexing transformation ops (linalg.transpose, tensor.collapse_shape,
     etc.) into a single iree_linalg_ext.map_scatter operation. For tensor.pad
     ops, the writing of pad values is distributed to workgroups and threads, and
     then the padding values are written directly to the output buffer of the
-    store_to_memref op.
+    store_to_buffer op.
   }];
   let dependentDialects = [
     "iree_compiler::IREE::LinalgExt::IREELinalgExtDialect",

compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_combine_layout_transformation.mlir

@@ -6,7 +6,7 @@ func.func @fold_pad_op(%source : tensor<250xf32>, %result : memref<256xf32>) {
   ^bb0(%arg0: index):
     tensor.yield %cst : f32
   } : tensor<250xf32> to tensor<256xf32>
-  iree_codegen.store_to_memref %padded, %result : tensor<256xf32> into memref<256xf32>
+  iree_codegen.store_to_buffer %padded, %result : tensor<256xf32> into memref<256xf32>
   return
 }
 //       CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (256, d0 + 64)>
@@ -25,7 +25,7 @@ func.func @fold_pad_op(%source : tensor<250xf32>, %result : memref<256xf32>) {
 //  CHECK-NEXT:   ^bb0(%[[IDX0:.+]]: index):
 //       CHECK:     iree_linalg_ext.yield %[[IDX0]], %[[TRUE]]
 //       CHECK:   } : tensor<250xf32> into tensor<256xf32> -> tensor<256xf32>
-//       CHECK:   iree_codegen.store_to_memref %[[MAP_SCATTER]], %[[RESULT]] : tensor<256xf32> into memref<256xf32>
+//       CHECK:   iree_codegen.store_to_buffer %[[MAP_SCATTER]], %[[RESULT]] : tensor<256xf32> into memref<256xf32>
 
 //       CHECK:   scf.forall (%[[WG_IV:.+]]) = (0) to (256) step (64) {
 //       CHECK:     %[[WG_TILE_UB:.+]] = affine.min #[[$MAP]](%[[WG_IV]])

compiler/src/iree/compiler/Codegen/Common/Passes.td

@@ -51,8 +51,8 @@ def BufferizeDispatchTensorLoadStorePass :
       "Bufferize the iree_tensor_ext.dispatch.tensor.load/store ops at dispatch boundaries";
   let description = [{
     Pass to bufferize the edges of dispatch regions, converting
-    iree_tensor_ext.dispatch.tensor.load ops to iree_codegen.load_from_memref, and
-    iree_tensor_ext.dispatch.tensor.store ops to iree_codegen.store_to_memref.
+    iree_tensor_ext.dispatch.tensor.load ops to iree_codegen.load_from_buffer, and
+    iree_tensor_ext.dispatch.tensor.store ops to iree_codegen.store_to_buffer.
   }];
   let dependentDialects = [
     "IREE::Codegen::IREECodegenDialect",
@@ -128,12 +128,12 @@ def CombineLayoutTransformationPass :
   let summary =
     "Combines layout transformation operations into a single map_scatter operation.";
   let description = [{
-    Starting from iree_codegen.store_to_memref ops, iteratively combine producer
+    Starting from iree_codegen.store_to_buffer ops, iteratively combine producer
     layout/indexing transformation ops (linalg.transpose, tensor.collapse_shape,
     etc.) into a single iree_linalg_ext.map_scatter operation. For tensor.pad
     ops, the writing of pad values is distributed to workgroups, and then the
     padding values are written directly to the output buffer of the
-    store_to_memref op.
+    store_to_buffer op.
   }];
   let dependentDialects = [
     "iree_compiler::IREE::LinalgExt::IREELinalgExtDialect",

compiler/src/iree/compiler/Codegen/Common/test/bufferize_dispatch_tensor_load_store.mlir

@@ -19,9 +19,9 @@ func.func @dispatch_tensor_load_and_store() {
 // CHECK-SAME:        binding(0) : memref<16xf32, #hal.descriptor_type<storage_buffer>>
 // CHECK:         %[[OUTPUT:.+]] = hal.interface.binding.subspan
 // CHECK-SAME:        binding(1) : memref<16xf32, #hal.descriptor_type<storage_buffer>>
-// CHECK:         %[[LOAD:.+]] = iree_codegen.load_from_memref %[[INPUT]]
+// CHECK:         %[[LOAD:.+]] = iree_codegen.load_from_buffer %[[INPUT]]
 // CHECK-SAME:        : memref<16xf32, #hal.descriptor_type<storage_buffer>> -> tensor<16xf32>
-// CHECK:         iree_codegen.store_to_memref %[[LOAD]], %[[OUTPUT]]
+// CHECK:         iree_codegen.store_to_buffer %[[LOAD]], %[[OUTPUT]]
 // CHECK-SAME:        : tensor<16xf32> into memref<16xf32, #hal.descriptor_type<storage_buffer>>
 
 // -----
@@ -51,9 +51,9 @@ func.func @dispatch_tensor_load_and_store_slices() {
 // CHECK:         %[[INPUT_SUBVIEW:.+]] = memref.subview %[[INPUT]][2] [12] [1]
 // CHECK-SAME:        : memref<16xf32, #hal.descriptor_type<storage_buffer>> to
 // CHECK-SAME:          memref<12xf32, strided<[1], offset: 2>, #hal.descriptor_type<storage_buffer>>
-// CHECK:         %[[LOAD:.+]] = iree_codegen.load_from_memref %[[INPUT_SUBVIEW]]
+// CHECK:         %[[LOAD:.+]] = iree_codegen.load_from_buffer %[[INPUT_SUBVIEW]]
 // CHECK-SAME:        : memref<12xf32, strided<[1], offset: 2>, #hal.descriptor_type<storage_buffer>> -> tensor<12xf32>
-// CHECK:         iree_codegen.store_to_memref %[[LOAD]], %[[OUTPUT_SUBVIEW]]
+// CHECK:         iree_codegen.store_to_buffer %[[LOAD]], %[[OUTPUT_SUBVIEW]]
 // CHECK-SAME:        : tensor<12xf32> into memref<12xf32, strided<[1], offset: 4>, #hal.descriptor_type<storage_buffer>>
 
 // -----
@@ -81,10 +81,10 @@ func.func @dispatch_tensor_load_and_store_with_compute_op() {
 // CHECK-SAME:        binding(1) : memref<16xf32, #hal.descriptor_type<storage_buffer>>
 // CHECK:         %[[OUTPUT_SUBVIEW:.+]] = memref.subview %[[OUTPUT]][4] [12] [1]
 // CHECK:         %[[INPUT_SUBVIEW:.+]] = memref.subview %[[INPUT]][2] [12] [1]
-// CHECK:         %[[LOAD:.+]] = iree_codegen.load_from_memref %[[INPUT_SUBVIEW]]
+// CHECK:         %[[LOAD:.+]] = iree_codegen.load_from_buffer %[[INPUT_SUBVIEW]]
 // CHECK:         %[[INIT:.+]] = tensor.empty() : tensor<12xf32>
 // CHECK:         %[[COPY:.+]] = linalg.copy ins(%[[LOAD]]{{.*}} outs(%[[INIT]]
-// CHECK:         iree_codegen.store_to_memref %[[COPY]], %[[OUTPUT_SUBVIEW]]
+// CHECK:         iree_codegen.store_to_buffer %[[COPY]], %[[OUTPUT_SUBVIEW]]
 
 // -----
 
@@ -117,9 +117,9 @@ func.func @dynamic_dispatch_tensor_load_and_store(%offset: index, %size: index,
 // CHECK:         %[[INPUT_SUBVIEW:.+]] = memref.subview %[[INPUT]][%[[OFFSET]]] [%[[SIZE]]] [%[[STRIDE]]]
 // CHECK-SAME:        : memref<?xf32, #hal.descriptor_type<storage_buffer>> to
 // CHECK-SAME:          memref<?xf32, strided<[?], offset: ?>, #hal.descriptor_type<storage_buffer>>
-// CHECK:         %[[LOAD:.+]] = iree_codegen.load_from_memref %[[INPUT_SUBVIEW]]
+// CHECK:         %[[LOAD:.+]] = iree_codegen.load_from_buffer %[[INPUT_SUBVIEW]]
 // CHECK-SAME:        : memref<?xf32, strided<[?], offset: ?>, #hal.descriptor_type<storage_buffer>> -> tensor<?xf32>
-// CHECK:         iree_codegen.store_to_memref %[[LOAD]], %[[OUTPUT_SUBVIEW]]
+// CHECK:         iree_codegen.store_to_buffer %[[LOAD]], %[[OUTPUT_SUBVIEW]]
 // CHECK-SAME:        : tensor<?xf32> into memref<?xf32, strided<[?], offset: ?>, #hal.descriptor_type<storage_buffer>>
 
 // -----
@@ -149,7 +149,7 @@ func.func @rank_reducing_slices() {
 // CHECK:         %[[INPUT_SUBVIEW:.+]] = memref.subview %[[INPUT]][0, 2] [1, 12] [1, 1]
 // CHECK-SAME:        : memref<8x16xf32, #hal.descriptor_type<storage_buffer>> to
 // CHECK-SAME:          memref<12xf32, strided<[1], offset: 2>, #hal.descriptor_type<storage_buffer>>
-// CHECK:         %[[LOAD:.+]] = iree_codegen.load_from_memref %[[INPUT_SUBVIEW]]
+// CHECK:         %[[LOAD:.+]] = iree_codegen.load_from_buffer %[[INPUT_SUBVIEW]]
 // CHECK-SAME:        : memref<12xf32, strided<[1], offset: 2>, #hal.descriptor_type<storage_buffer>> -> tensor<12xf32>
-// CHECK:         iree_codegen.store_to_memref %[[LOAD]], %[[OUTPUT_SUBVIEW]]
+// CHECK:         iree_codegen.store_to_buffer %[[LOAD]], %[[OUTPUT_SUBVIEW]]
 // CHECK-SAME:        : tensor<12xf32> into memref<12xf32, strided<[1], offset: 4>, #hal.descriptor_type<storage_buffer>>