Repository: iree-org/iree
Pull request #20589: [Codegen][GPU] Add placeholder op for buffer casts on tensors
Author: login: qedawkins
Created: 2025-04-21T19:56:07Z | Merged: 2025-04-30T16:42:13Z

Introduces iree_gpu.buffer_resource_cast to represent where to insert amdgpu.fat_raw_buffer_cast ops when bufferizing. It also supports taking a cache_swizzle stride. To avoid colliding the the existing plumbing of amdgpu buffer resources, there is a pass that walks producers of these cast ops and either drops the casts or drops the annotation on the binding.

Right now the verification for cache_swizzle is simply checking for a chain of single-use producers, but in the future we can add more involved verification that checks that all users specify the same cache swizzle value.


Diff:
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel
index 50dbbe7278b8..5266b6a515ed 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel
@@ -53,6 +53,7 @@ iree_compiler_cc_library(
         "DecomposeHorizontallyFusedGemms.cpp",
         "ExpandGPUOps.cpp",
         "GPUApplyTilingLevel.cpp",
+        "GPUBubbleResourceCasts.cpp",
         "GPUCheckResourceUsage.cpp",
         "GPUCombineValueBarriers.cpp",
         "GPUCreateFastSlowPath.cpp",
@@ -109,6 +110,7 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Codegen/Utils:VectorOpUtils",
         "//compiler/src/iree/compiler/Dialect/HAL/IR",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/Utils",
+        "//compiler/src/iree/compiler/Dialect/TensorExt/IR",
         "//compiler/src/iree/compiler/Utils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AMDGPUDialect",
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt
index b61d4b4c57f5..f3df8ae717e3 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt
@@ -51,6 +51,7 @@ iree_cc_library(
     "DecomposeHorizontallyFusedGemms.cpp"
     "ExpandGPUOps.cpp"
     "GPUApplyTilingLevel.cpp"
+    "GPUBubbleResourceCasts.cpp"
     "GPUCheckResourceUsage.cpp"
     "GPUCombineValueBarriers.cpp"
     "GPUCreateFastSlowPath.cpp"
@@ -143,6 +144,7 @@ iree_cc_library(
     iree::compiler::Codegen::Utils::VectorOpUtils
     iree::compiler::Dialect::HAL::IR
     iree::compiler::Dialect::LinalgExt::Utils
+    iree::compiler::Dialect::TensorExt::IR
     iree::compiler::Utils
   PUBLIC
 )
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUBubbleResourceCasts.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUBubbleResourceCasts.cpp
new file mode 100644
index 000000000000..b27b4863dcbb
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUBubbleResourceCasts.cpp
@@ -0,0 +1,201 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Codegen/Common/GPU/Passes.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
+#include "iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler {
+
+#define GEN_PASS_DEF_GPUBUBBLERESOURCECASTSPASS
+#include "iree/compiler/Codegen/Common/GPU/Passes.h.inc"
+
+namespace {
+struct GPUBubbleResourceCastsPass final
+    : impl::GPUBubbleResourceCastsPassBase<GPUBubbleResourceCastsPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+/// Walk the producers of |input| and return the hal.interface.binding_subspan
+/// that produces it if the binding is readonly.
+static Operation *isTriviallyProducedByReadOnlyViewLike(Value input) {
+  Value next = input;
+  while (auto slice = next.getDefiningOp<tensor::ExtractSliceOp>()) {
+    next = slice.getSource();
+  }
+
+  auto tensorLoad = next.getDefiningOp<IREE::TensorExt::DispatchTensorLoadOp>();
+  if (!tensorLoad) {
+    return nullptr;
+  }
+
+  auto binding = tensorLoad.getSource()
+                     .getDefiningOp<IREE::HAL::InterfaceBindingSubspanOp>();
+  if (!binding ||
+      cast<IREE::TensorExt::DispatchTensorType>(binding.getType())
+              .getAccess() != IREE::TensorExt::TensorAccess::ReadOnly) {
+    return nullptr;
+  }
+  return binding;
+}
+
+namespace {
+struct BubbleResourceCastPattern
+    : public OpRewritePattern<IREE::GPU::BufferResourceCastOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(IREE::GPU::BufferResourceCastOp castOp,
+                                PatternRewriter &rewriter) const override {
+    // Skip ops with cache swizzle strides because the stride could be
+    // rendered invalid by the swap.
+    // TODO: Support bubbling swizzled casts in cases where it makes sense.
+    // e.g. iree_gpu.buffer_resource_cast(tensor.extract_slice) if the cache
+    // swizzle stride does not depend on the extract_slice.
+    if (castOp.getCacheSwizzleStride()) {
+      return failure();
+    }
+
+    auto producer = castOp.getInput().getDefiningOp();
+    if (!producer) {
+      return failure();
+    }
+
+    Location loc = castOp.getLoc();
+    bool swapped =
+        TypeSwitch<Operation *, bool>(producer)
+            .Case<tensor::ExtractSliceOp>([&](tensor::ExtractSliceOp extract) {
+              if (!extract->hasOneUse()) {
+                return false;
+              }
+
+              rewriter.setInsertionPoint(extract);
+              auto newCast = rewriter.create<IREE::GPU::BufferResourceCastOp>(
+                  loc, extract.getSource().getType(), extract.getSource());
+              extract.getSourceMutable().assign(newCast);
+              return true;
+            })
+            .Case<tensor::ExpandShapeOp>([&](tensor::ExpandShapeOp expand) {
+              if (!expand->hasOneUse()) {
+                return false;
+              }
+
+              rewriter.setInsertionPoint(expand);
+              auto newCast = rewriter.create<IREE::GPU::BufferResourceCastOp>(
+                  loc, expand.getSrcType(), expand.getSrc());
+              expand.getSrcMutable().assign(newCast);
+              return true;
+            })
+            .Case<tensor::CollapseShapeOp>(
+                [&](tensor::CollapseShapeOp collapse) {
+                  if (!collapse->hasOneUse()) {
+                    return false;
+                  }
+
+                  rewriter.setInsertionPoint(collapse);
+                  auto newCast =
+                      rewriter.create<IREE::GPU::BufferResourceCastOp>(
+                          loc, collapse.getSrcType(), collapse.getSrc());
+                  collapse.getSrcMutable().assign(newCast);
+                  return true;
+                })
+            .Case<tensor::PadOp>([&](tensor::PadOp pad) {
+              if (!pad->hasOneUse()) {
+                return false;
+              }
+
+              rewriter.setInsertionPoint(pad);
+              auto newCast = rewriter.create<IREE::GPU::BufferResourceCastOp>(
+                  loc, pad.getSourceType(), pad.getSource());
+              pad.getSourceMutable().assign(newCast);
+              return true;
+            })
+            .Case<linalg::LinalgOp>([&](linalg::LinalgOp linalgOp) {
+              // Skip gather-like linalg ops.
+              if (linalgOp.hasIndexSemantics() || !linalgOp->hasOneUse()) {
+                return false;
+              }
+              rewriter.setInsertionPoint(linalgOp);
+              // Only propagate to input operands.
+              for (auto inputOperand : linalgOp.getDpsInputOperands()) {
+                auto newCast = rewriter.create<IREE::GPU::BufferResourceCastOp>(
+                    loc, inputOperand->get().getType(), inputOperand->get());
+                inputOperand->assign(newCast);
+              }
+              return true;
+            })
+            .Case<tensor::EmptyOp>([&](tensor::EmptyOp empty) {
+              // Drop all casts of empties.
+              return true;
+            })
+            .Default([](Operation *op) {
+              // Leave all other casts alone.
+              return false;
+            });
+
+    // If swapping did not occur, do not drop the cast.
+    if (!swapped) {
+      return failure();
+    }
+
+    rewriter.replaceOp(castOp, castOp.getInput());
+    return success();
+  }
+};
+} // namespace
+
+void GPUBubbleResourceCastsPass::runOnOperation() {
+  MLIRContext *context = &getContext();
+  Operation *op = getOperation();
+
+  {
+    MLIRContext *context = &getContext();
+    RewritePatternSet patterns(context);
+    patterns.add<BubbleResourceCastPattern>(context);
+    if (failed(applyPatternsGreedily(op, std::move(patterns)))) {
+      return signalPassFailure();
+    }
+  }
+
+  // GPU Dialect to lookup rocdl. This pass depends on IREEGPUDialect so this
+  // always succeeds.
+  auto *ireeGpuDialect = context->getLoadedDialect<IREE::GPU::IREEGPUDialect>();
+
+  SmallVector<IREE::GPU::BufferResourceCastOp> castsToDrop;
+  op->walk([&](IREE::GPU::BufferResourceCastOp castOp) {
+    // Skip ops without cache swizzle values set. There is no need to drop
+    // these because we will try to bubble them up.
+    if (!castOp.getCacheSwizzleStride()) {
+      return;
+    }
+    if (Operation *binding =
+            isTriviallyProducedByReadOnlyViewLike(castOp.getInput())) {
+      if (ireeGpuDialect->getUseRocdlBufferInstructionsAttrHelper()
+              .isAttrPresent(binding)) {
+        // TODO: Support updating buffer structs so we don't have to replace
+        // the cast on the binding with this cast.
+        ireeGpuDialect->getUseRocdlBufferInstructionsAttrHelper().removeAttr(
+            binding);
+      }
+    } else {
+      // If the cast can't trivially map to a view-like producer of a binding,
+      // drop it. Common if there are any unfused producers.
+      castsToDrop.push_back(castOp);
+    }
+  });
+
+  // Drop all cache swizzling ops that might be blocking fusion.
+  for (IREE::GPU::BufferResourceCastOp cast : castsToDrop) {
+    cast->replaceAllUsesWith(ValueRange{cast.getInput()});
+  }
+}
+
+}; // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td b/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
index 4a872fa48357..66fc651e2866 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
@@ -90,6 +90,14 @@ def GPUDistributeScfForPass :
   ];
 }
 
+def GPUBubbleResourceCastsPass :
+    Pass<"iree-codegen-gpu-bubble-resource-casts", ""> {
+  let summary = "Bubbles iree_gpu.buffer_resource_cast ops upwards.";
+  let dependentDialects = [
+    "::mlir::iree_compiler::IREE::GPU::IREEGPUDialect"
+  ];
+}
+
 def GPUFuseAndHoistParallelLoopsPass :
     InterfacePass<"iree-codegen-gpu-fuse-and-hoist-parallel-loops", "mlir::FunctionOpInterface"> {
   let summary = "Greedily fuses and hoists parallel loops.";
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel
index 4d3957c951b8..b9dfdd31fbcd 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel
@@ -21,6 +21,7 @@ iree_lit_test_suite(
             "decompose_horizontally_fused_gemms.mlir",
             "gpu_apply_derived_thread_config.mlir",
             "gpu_apply_tiling_level.mlir",
+            "gpu_bubble_resource_casts.mlir",
             "gpu_check_resource_usage.mlir",
             "gpu_create_fast_slow_path.mlir",
             "gpu_distribute.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt
index bd92c906b6aa..724fd18752ac 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt
@@ -17,6 +17,7 @@ iree_lit_test_suite(
     "decompose_horizontally_fused_gemms.mlir"
     "gpu_apply_derived_thread_config.mlir"
     "gpu_apply_tiling_level.mlir"
+    "gpu_bubble_resource_casts.mlir"
     "gpu_check_resource_usage.mlir"
     "gpu_combine_value_barriers.mlir"
     "gpu_create_fast_slow_path.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_bubble_resource_casts.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_bubble_resource_casts.mlir
new file mode 100644
index 000000000000..abe8d06b9f0f
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_bubble_resource_casts.mlir
@@ -0,0 +1,128 @@
+// RUN: iree-opt %s --split-input-file --pass-pipeline="builtin.module(func.func(iree-codegen-gpu-bubble-resource-casts))" | FileCheck %s
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [#hal.pipeline.binding<storage_buffer>]>
+func.func @simple_cast() -> (tensor<2xf32>, tensor<2xf32>) {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %arg0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>>
+  %0 = iree_tensor_ext.dispatch.tensor.load %arg0, offsets=[0], sizes=[2], strides=[1]
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>> -> tensor<2xf32>
+  %1 = iree_gpu.buffer_resource_cast %0 : tensor<2xf32>
+  return %0, %1 : tensor<2xf32>, tensor<2xf32>
+}
+
+// CHECK-LABEL: func.func @simple_cast
+//       CHECK:   %[[LOAD:.+]] = iree_tensor_ext.dispatch.tensor.load
+//       CHECK:   %[[CAST:.+]] = iree_gpu.buffer_resource_cast
+//       CHECK:   return %[[LOAD]], %[[CAST]] : tensor<2xf32>, tensor<2xf32>
+
+// -----
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [#hal.pipeline.binding<storage_buffer>]>
+func.func @drop_binding_annotation() -> tensor<2xf32> {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %arg0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+    {iree_gpu.use_rocdl_buffer_instructions}
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>>
+  %0 = iree_tensor_ext.dispatch.tensor.load %arg0, offsets=[0], sizes=[2], strides=[1]
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>> -> tensor<2xf32>
+  %1 = iree_gpu.buffer_resource_cast %0 cacheSwizzleStride(%c2) : tensor<2xf32>
+  return %1 : tensor<2xf32>
+}
+
+// CHECK-LABEL: func.func @drop_binding_annotation
+//       CHECK:   hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%{{.*}}) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>>
+//       CHECK:   %[[CAST:.+]] = iree_gpu.buffer_resource_cast
+//       CHECK:   return %[[CAST]] : tensor<2xf32>
+
+// -----
+
+func.func @bubble_extract_slice(%arg0: tensor<4xf32>) -> tensor<2xf32> {
+  %0 = tensor.extract_slice %arg0 [0] [2] [1] : tensor<4xf32> to tensor<2xf32>
+  %1 = iree_gpu.buffer_resource_cast %0 : tensor<2xf32>
+  return %1 : tensor<2xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_extract_slice
+//  CHECK-SAME:   %[[ARG0:[A-Za-z0-9]+]]: tensor<4xf32>
+//       CHECK:   %[[CAST:.+]] = iree_gpu.buffer_resource_cast %[[ARG0]]
+//       CHECK:   %[[EXTRACT:.+]] = tensor.extract_slice %[[CAST]]
+//       CHECK:   return %[[EXTRACT]]
+
+// -----
+
+func.func @bubble_pad(%arg0: tensor<1xf32>) -> tensor<2xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.pad %arg0 low[0] high[1] {
+  ^bb0(%arg1: index):
+    tensor.yield %cst : f32
+  } : tensor<1xf32> to tensor<2xf32>
+  %1 = iree_gpu.buffer_resource_cast %0 : tensor<2xf32>
+  return %1 : tensor<2xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_pad
+//  CHECK-SAME:   %[[ARG0:[A-Za-z0-9]+]]: tensor<1xf32>
+//       CHECK:   %[[CAST:.+]] = iree_gpu.buffer_resource_cast %[[ARG0]]
+//       CHECK:   %[[PAD:.+]] = tensor.pad %[[CAST]]
+//       CHECK:   return %[[PAD]]
+
+// -----
+
+func.func @bubble_expand_and_collapse(%arg0: tensor<3x2xf32>) -> tensor<2x3xf32> {
+  %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x2xf32> into tensor<6xf32>
+  %1 = tensor.expand_shape %0 [[0, 1]] output_shape [3, 2]
+      : tensor<6xf32> into tensor<2x3xf32>
+  %2 = iree_gpu.buffer_resource_cast %1 : tensor<2x3xf32>
+  return %2 : tensor<2x3xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_expand_and_collapse
+//  CHECK-SAME:   %[[ARG0:[A-Za-z0-9]+]]: tensor<3x2xf32>
+//       CHECK:   %[[CAST:.+]] = iree_gpu.buffer_resource_cast %[[ARG0]]
+//       CHECK:   %[[COLLAPSE:.+]] = tensor.collapse_shape %[[CAST]]
+//       CHECK:   %[[EXPAND:.+]] = tensor.expand_shape %[[COLLAPSE]]
+//       CHECK:   return %[[EXPAND]]
+
+// -----
+
+func.func @drop_blocked_resource_cast(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.pad %arg0 low[0] high[1] {
+  ^bb0(%arg1: index):
+    tensor.yield %cst : f32
+  } : tensor<?xf32> to tensor<?xf32>
+  %swizzle = tensor.dim %arg0, %c0 : tensor<?xf32>
+  %1 = iree_gpu.buffer_resource_cast %0 cacheSwizzleStride(%swizzle) : tensor<?xf32>
+  return %1 : tensor<?xf32>
+}
+
+// CHECK-LABEL: func.func @drop_blocked_resource_cast
+//  CHECK-SAME:   %[[ARG0:[A-Za-z0-9]+]]: tensor<?xf32>
+//       CHECK:   %[[PAD:.+]] = tensor.pad %[[ARG0]]
+//       CHECK:   return %[[PAD]]
+
+// -----
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [#hal.pipeline.binding<storage_buffer>]>
+func.func @keep_swizzle_cast() -> tensor<1xf32> {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %arg0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>>
+  %0 = iree_tensor_ext.dispatch.tensor.load %arg0, offsets=[0], sizes=[2], strides=[1]
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>> -> tensor<2xf32>
+  %1 = tensor.extract_slice %0 [0] [1] [1] : tensor<2xf32> to tensor<1xf32>
+  %2 = iree_gpu.buffer_resource_cast %1 cacheSwizzleStride(%c2) : tensor<1xf32>
+  return %2 : tensor<1xf32>
+}
+
+// CHECK-LABEL: func.func @keep_swizzle_cast
+//       CHECK:   %[[BINDING:.+]] = hal.interface.binding.subspan
+//       CHECK:   %[[LOAD:.+]] = iree_tensor_ext.dispatch.tensor.load %[[BINDING]]
+//       CHECK:   %[[EXTRACT:.+]] = tensor.extract_slice %[[LOAD]]
+//       CHECK:   %[[CAST:.+]] = iree_gpu.buffer_resource_cast %[[EXTRACT]]
+//       CHECK:   return %[[CAST]] : tensor<1xf32>
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/iree_comprehensive_bufferize.mlir b/compiler/src/iree/compiler/Codegen/Common/test/iree_comprehensive_bufferize.mlir
index 1e46a1e83480..d81ec0646500 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/iree_comprehensive_bufferize.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/iree_comprehensive_bufferize.mlir
@@ -3031,3 +3031,82 @@ func.func @load_from_memref_read_only_copy() {
 //       CHECK:   linalg.generic
 //  CHECK-SAME:       ins(%[[ALLOC]] :
 //  CHECK-SAME:       outs(%[[OUTPUT]] :
+
+// -----
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [
+  #hal.pipeline.binding<storage_buffer>,
+  #hal.pipeline.binding<storage_buffer>
+]>
+func.func @simple_resource_cast() {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %arg0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>>
+  %arg1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0)
+    : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<2xf32>>
+  %0 = iree_tensor_ext.dispatch.tensor.load %arg0, offsets=[0], sizes=[2], strides=[1]
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>> -> tensor<2xf32>
+  %1 = iree_gpu.buffer_resource_cast %0 : tensor<2xf32>
+  iree_tensor_ext.dispatch.tensor.store %1, %arg1, offsets=[0], sizes=[2], strides=[1]
+    : tensor<2xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<2xf32>>
+  return
+}
+
+// CHECK-LABEL: func.func @simple_resource_cast
+//       CHECK:   %[[INPUT:.+]] = hal.interface.binding.subspan layout({{.*}}) binding(0)
+//       CHECK:   %[[CAST:.+]] = amdgpu.fat_raw_buffer_cast %[[INPUT]] resetOffset
+//  CHECK-SAME:     memref<2xf32, #hal.descriptor_type<storage_buffer>> to memref<2xf32, #amdgpu.address_space<fat_raw_buffer>>
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     ins(%[[CAST]] : memref<2xf32, #amdgpu.address_space<fat_raw_buffer>>)
+//  CHECK-SAME:     outs(%{{.*}} : memref<2xf32, #hal.descriptor_type<storage_buffer>>)
+
+// -----
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [
+  #hal.pipeline.binding<storage_buffer>
+]>
+func.func @drop_resource_cast_if_not_storage_buffer() -> tensor<2xf32> {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %arg0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>>
+  %0 = iree_tensor_ext.dispatch.tensor.load %arg0, offsets=[0], sizes=[2], strides=[1]
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>> -> tensor<2xf32>
+  %empty = tensor.empty() : tensor<2xf32>
+  %copy = linalg.copy ins(%0 : tensor<2xf32>) outs(%empty : tensor<2xf32>) -> tensor<2xf32>
+  %1 = iree_gpu.buffer_resource_cast %copy : tensor<2xf32>
+  return %1 : tensor<2xf32>
+}
+
+// CHECK-LABEL: func.func @drop_resource_cast_if_not_storage_buffer
+//   CHECK-NOT:   amdgpu.fat_raw_buffer_cast
+
+// -----
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [
+  #hal.pipeline.binding<storage_buffer>,
+  #hal.pipeline.binding<storage_buffer>
+]>
+func.func @cache_swizzle_resource_cast(%stride: index) {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %arg0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>>
+  %arg1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0)
+    : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<2xf32>>
+  %0 = iree_tensor_ext.dispatch.tensor.load %arg0, offsets=[0], sizes=[2], strides=[1]
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2xf32>> -> tensor<2xf32>
+  %1 = iree_gpu.buffer_resource_cast %0 cacheSwizzleStride(%stride) : tensor<2xf32>
+  iree_tensor_ext.dispatch.tensor.store %1, %arg1, offsets=[0], sizes=[2], strides=[1]
+    : tensor<2xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<2xf32>>
+  return
+}
+
+// CHECK-LABEL: func.func @cache_swizzle_resource_cast
+//  CHECK-SAME:   %[[STRIDE:[A-Za-z0-9]+]]: index
+//   CHECK-DAG:   %[[INPUT:.+]] = hal.interface.binding.subspan layout({{.*}}) binding(0)
+//   CHECK-DAG:   %[[TRUNC:.+]] = arith.index_cast %[[STRIDE]] : index to i14
+//       CHECK:   %[[CAST:.+]] = amdgpu.fat_raw_buffer_cast %[[INPUT]]
+//  CHECK-SAME:     cacheSwizzleStride(%[[TRUNC]]) resetOffset
+//  CHECK-SAME:     memref<2xf32, #hal.descriptor_type<storage_buffer>> to memref<2xf32, #amdgpu.address_space<fat_raw_buffer>>
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.cpp
index ff3535305682..753fcb3cd47c 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.cpp
@@ -370,4 +370,79 @@ LogicalResult ValueBarrierOp::verify() {
   return success();
 }
 
+// AMD Specific Operations
+
+//===----------------------------------------------------------------------===//
+// BufferResourceCastOp
+//===----------------------------------------------------------------------===//
+
+static RankedTensorType getMaximumStaticType(tensor::CastOp castOp) {
+  auto inputType = dyn_cast<RankedTensorType>(castOp.getSource().getType());
+  auto resultType = dyn_cast<RankedTensorType>(castOp.getType());
+  if (!inputType || !resultType) {
+    return RankedTensorType();
+  }
+
+  assert(inputType.getRank() == resultType.getRank() &&
+         "Rank must match for ranked -> ranked cast");
+
+  SmallVector<int64_t> join;
+  join.reserve(inputType.getRank());
+  for (int64_t i = 0; i < inputType.getRank(); ++i) {
+    if (inputType.isDynamicDim(i)) {
+      join.push_back(resultType.getDimSize(i));
+      continue;
+    }
+    if (resultType.isDynamicDim(i)) {
+      join.push_back(inputType.getDimSize(i));
+      continue;
+    }
+
+    // Cast verifier requires that static sizes match.
+    join.push_back(inputType.getDimSize(i));
+  }
+  return RankedTensorType::get(join, inputType.getElementType());
+}
+
+struct FoldBufferCastOfTensorCast final
+    : OpRewritePattern<BufferResourceCastOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(BufferResourceCastOp castOp,
+                                PatternRewriter &rewriter) const override {
+    // Check whether the cast increases the amount of available static info.
+    auto tensorCast = castOp.getInput().getDefiningOp<tensor::CastOp>();
+    if (!tensorCast) {
+      return failure();
+    }
+
+    RankedTensorType maxStaticType = getMaximumStaticType(tensorCast);
+    if (!maxStaticType || maxStaticType == castOp.getInput().getType()) {
+      return failure();
+    }
+
+    Value newSource = tensorCast.getSource();
+    if (newSource.getType() != maxStaticType) {
+      // Cast to the type with maximum static information if the input and
+      // result types contain different static info.
+      newSource = rewriter.create<tensor::CastOp>(castOp.getLoc(),
+                                                  maxStaticType, newSource);
+    }
+    auto newBufferCast = rewriter.create<IREE::GPU::BufferResourceCastOp>(
+        castOp.getLoc(), maxStaticType, newSource,
+        castOp.getCacheSwizzleStride());
+    newBufferCast->setDiscardableAttrs(castOp->getDiscardableAttrDictionary());
+
+    // Cast back to the original result type.
+    rewriter.replaceOpWithNewOp<tensor::CastOp>(
+        castOp, castOp.getResult().getType(), newBufferCast);
+    return success();
+  };
+};
+
+void BufferResourceCastOp::getCanonicalizationPatterns(
+    RewritePatternSet &results, MLIRContext *ctx) {
+  results.add<FoldBufferCastOfTensorCast>(ctx);
+}
+
 } // namespace mlir::iree_compiler::IREE::GPU
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.td b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.td
index 18ae9c357bce..a4c9e7608fa9 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.td
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.td
@@ -456,4 +456,44 @@ def IREEGPU_YieldOp : Op<IREEGPU_Dialect, "yield", [
       [{  attr-dict ($values^ `:` type($values))? }];
 }
 
+//===----------------------------------------------------------------------===//
+//
+// AMD Specific Operations
+//
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// BufferResourceCastOp
+//===----------------------------------------------------------------------===//
+
+def IREEGPU_BufferResourceCastOp : Op<IREEGPU_Dialect, "buffer_resource_cast", [
+  Pure,
+  AllTypesMatch<["input", "result"]>]> {
+  let summary = "Represents a cast to addr_space<7> (buffer resource) before bufferization";
+  let description = [{
+    Nominal cast of a tensor to AMDGPU buffer resource memory space before
+    bufferization. This op takes the parameters with which to perform the cast
+    if |input| bufferizes to `storage_buffer` memory space. If |input| resolves
+    to any other memory space this op is silently dropped and has no effect.
+
+    If |cache_swizzle_stride| is present, there is verification before
+    bufferization that all producers of |input| are view-like and single source
+    and user (i.e. trivially no alias). In all other cases this op is best
+    effort and has no verification or failure modes.
+
+    // TODO: Add other parameters for casting as needed.
+  }];
+
+  let arguments = (ins  AnyRankedTensor:$input,
+                        Optional<Index>:$cache_swizzle_stride);
+  let results   = (outs AnyRankedTensor:$result);
+
+  let assemblyFormat = [{
+    $input oilist (`cacheSwizzleStride` `(` $cache_swizzle_stride `)` )
+    attr-dict `:` type($result)
+  }];
+
+  let hasCanonicalizer = 1;
+}
+
 #endif // IREE_CODEGEN_DIALECT_IREEGPUOPS
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/BUILD.bazel
index d477831f5aab..77bd72dcb20d 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/BUILD.bazel
@@ -18,6 +18,7 @@ iree_lit_test_suite(
     name = "lit",
     srcs = enforce_glob(
         [
+            "canonicalize.mlir",
             "iree_gpu_attrs.mlir",
             "iree_gpu_ops.mlir",
             "target_attrs.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/CMakeLists.txt
index a7085c6ca1b1..6217a5f18e11 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/CMakeLists.txt
@@ -14,6 +14,7 @@ iree_lit_test_suite(
   NAME
     lit
   SRCS
+    "canonicalize.mlir"
     "iree_gpu_attrs.mlir"
     "iree_gpu_ops.mlir"
     "target_attrs.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/canonicalize.mlir b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/canonicalize.mlir
new file mode 100644
index 000000000000..ade137afb0a8
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/canonicalize.mlir
@@ -0,0 +1,42 @@
+// RUN: iree-opt %s --canonicalize --split-input-file | FileCheck %s
+
+func.func @fold_cast_into_resource_cast(%input: tensor<32xf16>) -> tensor<?xf16> {
+  %cast = tensor.cast %input : tensor<32xf16> to tensor<?xf16>
+  %resource_cast = iree_gpu.buffer_resource_cast %cast : tensor<?xf16>
+  return %resource_cast : tensor<?xf16>
+}
+
+// CHECK-LABEL: func @fold_cast_into_resource_cast
+//  CHECK-SAME:   %[[INPUT:[A-Za-z0-9]+]]: tensor<32xf16>
+//       CHECK:   %[[RESOURCE_CAST:.+]] = iree_gpu.buffer_resource_cast %[[INPUT]] : tensor<32xf16>
+//       CHECK:   %[[CAST:.+]] = tensor.cast %[[RESOURCE_CAST]] : tensor<32xf16> to tensor<?xf16>
+//       CHECK:   return %[[CAST]]
+
+// -----
+
+func.func @no_fold_static_cast_into_resource_cast(%input: tensor<?xf16>) -> tensor<32xf16> {
+  %cast = tensor.cast %input : tensor<?xf16> to tensor<32xf16>
+  %resource_cast = iree_gpu.buffer_resource_cast %cast : tensor<32xf16>
+  return %resource_cast : tensor<32xf16>
+}
+
+// CHECK-LABEL: func @no_fold_static_cast_into_resource_cast
+//  CHECK-SAME:   %[[INPUT:[A-Za-z0-9]+]]: tensor<?xf16>
+//       CHECK:   %[[CAST:.+]] = tensor.cast %[[INPUT]] : tensor<?xf16> to tensor<32xf16>
+//       CHECK:   %[[RESOURCE_CAST:.+]] = iree_gpu.buffer_resource_cast %[[CAST]] : tensor<32xf16>
+//       CHECK:   return %[[RESOURCE_CAST]]
+
+// -----
+
+func.func @split_cast_adding_and_removing_static_info(%input: tensor<?x32xf16>) -> tensor<32x?xf16> {
+  %cast = tensor.cast %input : tensor<?x32xf16> to tensor<32x?xf16>
+  %resource_cast = iree_gpu.buffer_resource_cast %cast : tensor<32x?xf16>
+  return %resource_cast : tensor<32x?xf16>
+}
+
+// CHECK-LABEL: func @split_cast_adding_and_removing_static_info
+//  CHECK-SAME:   %[[INPUT:[A-Za-z0-9]+]]: tensor<?x32xf16>
+//       CHECK:   %[[CAST_IN:.+]] = tensor.cast %[[INPUT]] : tensor<?x32xf16> to tensor<32x32xf16>
+//       CHECK:   %[[RESOURCE_CAST:.+]] = iree_gpu.buffer_resource_cast %[[CAST_IN]] : tensor<32x32xf16>
+//       CHECK:   %[[CAST_OUT:.+]] = tensor.cast %[[RESOURCE_CAST]] : tensor<32x32xf16> to tensor<32x?xf16>
+//       CHECK:   return %[[CAST_OUT]]
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/BUILD.bazel
index b98d6dd03d87..61e74f50736a 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/BUILD.bazel
@@ -22,6 +22,8 @@ iree_compiler_cc_library(
     ],
     deps = [
         "//compiler/src/iree/compiler/Codegen/Dialect/GPU/IR:IREEGPUDialect",
+        "//compiler/src/iree/compiler/Dialect/HAL/IR",
+        "@llvm-project//mlir:AMDGPUDialect",
         "@llvm-project//mlir:BufferizationDialect",
         "@llvm-project//mlir:BufferizationInterfaces",
         "@llvm-project//mlir:BufferizationTransforms",
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/BufferizationInterfaces.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/BufferizationInterfaces.cpp
index 17202adc57be..332e287b342f 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/BufferizationInterfaces.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/BufferizationInterfaces.cpp
@@ -8,6 +8,8 @@
 
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.h"
+#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
+#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
@@ -161,7 +163,7 @@ struct BarrierRegionOpBufferizationInterface
   }
 };
 
-/// Bufferization of iree_gpu.tensor_barrier. Always just bufferizes in place
+/// Bufferization of iree_gpu.value_barrier. Always just bufferizes in place
 /// and replaces with a barrier.
 struct ValueBarrierOpBufferizationInterface
     : public BufferizableOpInterface::ExternalModel<
@@ -281,6 +283,111 @@ struct YieldOpBufferizationInterface
   }
 };
 
+/// AMD Specific Ops
+
+static bool hasStorageBufferMemSpace(BaseMemRefType m) {
+  Attribute maybeMemorySpace = m.getMemorySpace();
+  Builder b(m.getContext());
+  Attribute storageBufferMemSpace = b.getAttr<IREE::HAL::DescriptorTypeAttr>(
+      IREE::HAL::DescriptorType::StorageBuffer);
+  return maybeMemorySpace && maybeMemorySpace == storageBufferMemSpace;
+}
+
+/// Bufferization of iree_gpu.buffer_resource_cast. Bufferizes to
+/// amdgpu.fat_raw_buffer_cast if the source memref is of memory space
+/// `storage_buffer`, else just forwards the input. This op never
+/// reads or writes.
+struct BufferResourceCastOpBufferizationInterface
+    : public BufferizableOpInterface::ExternalModel<
+          BufferResourceCastOpBufferizationInterface,
+          IREE::GPU::BufferResourceCastOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    // This op never needs to bufferize to a copy.
+    return false;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    return false;
+  }
+
+  bufferization::AliasingValueList
+  getAliasingValues(Operation *op, OpOperand &opOperand,
+                    const AnalysisState &state) const {
+    auto castOp = cast<IREE::GPU::BufferResourceCastOp>(op);
+    SmallVector<bufferization::AliasingValue> alist;
+    if (opOperand.get() == castOp.getInput()) {
+      alist.push_back({castOp.getResult(), BufferRelation::Equivalent});
+    }
+    return alist;
+  }
+
+  FailureOr<BaseMemRefType>
+  getBufferType(Operation *op, Value value, const BufferizationOptions &options,
+                SmallVector<Value> &invocationStack) const {
+    auto castOp = cast<IREE::GPU::BufferResourceCastOp>(op);
+    assert(value.getDefiningOp() == castOp && "invalid value");
+    auto srcMemrefType = bufferization::getBufferType(castOp.getInput(),
+                                                      options, invocationStack);
+    if (failed(srcMemrefType))
+      return failure();
+
+    if (!hasStorageBufferMemSpace(srcMemrefType.value())) {
+      return srcMemrefType;
+    }
+
+    auto rankedSrcType = cast<MemRefType>(srcMemrefType.value());
+
+    // Sad AMDGPU dep.
+    Attribute bufferMemSpace = amdgpu::AddressSpaceAttr::get(
+        op->getContext(), amdgpu::AddressSpace::FatRawBuffer);
+    return MemRefType::get(rankedSrcType.getShape(),
+                           rankedSrcType.getElementType(),
+                           rankedSrcType.getLayout(), bufferMemSpace);
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationOptions &options) const {
+    auto castOp = cast<IREE::GPU::BufferResourceCastOp>(op);
+
+    FailureOr<Value> buffer = getBuffer(rewriter, castOp.getInput(), options);
+    if (failed(buffer)) {
+      return failure();
+    }
+
+    // This operation either bufferizes in place or with a cast.
+    if (hasStorageBufferMemSpace(cast<MemRefType>(buffer.value().getType()))) {
+      Location loc = castOp.getLoc();
+      Value cacheSwizzleStride = Value{};
+      if (auto maybeIndexCacheSwizzle = castOp.getCacheSwizzleStride()) {
+        // Cache swizzle supports only upto 8k stride. Also simply swizzling the
+        // largest available stride (8k) doesn't help those unsupported large
+        // stride. Especially better to avoid using the stride which is 2^N when
+        // N>13, e.g. by add padding to the buffer.
+        //
+        // stride[13:0] = swizzling stride
+        // stride[14] = swizzle enabling bit
+        // FatRawBufferCast's lowering handles this for us. Just truncate to 14
+        // bits.
+        Type i14Type = rewriter.getIntegerType(14);
+        cacheSwizzleStride = rewriter.create<arith::IndexCastOp>(
+            loc, i14Type, maybeIndexCacheSwizzle);
+      }
+      buffer = rewriter
+                   .create<amdgpu::FatRawBufferCastOp>(
+                       loc, buffer.value(), /*validBytes=*/Value{},
+                       /*cacheSwizzleStride=*/cacheSwizzleStride,
+                       /*boundsCheck=*/true,
+                       /*resetOffset=*/true)
+                   .getResult();
+    }
+
+    bufferization::replaceOpWithBufferizedValues(rewriter, op, buffer.value());
+    return success();
+  }
+};
+
 } // namespace
 
 void registerIREEGPUBufferizationInterfaces(DialectRegistry &registry) {
@@ -292,6 +399,9 @@ void registerIREEGPUBufferizationInterfaces(DialectRegistry &registry) {
             ValueBarrierOpBufferizationInterface>(*context);
         IREE::GPU::YieldOp::attachInterface<YieldOpBufferizationInterface>(
             *context);
+
+        IREE::GPU::BufferResourceCastOp::attachInterface<
+            BufferResourceCastOpBufferizationInterface>(*context);
       });
 }
 
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/CMakeLists.txt
index 23a6aa9eebcb..df6792ba2392 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/CMakeLists.txt
@@ -18,6 +18,7 @@ iree_cc_library(
   SRCS
     "BufferizationInterfaces.cpp"
   DEPS
+    MLIRAMDGPUDialect
     MLIRBufferizationDialect
     MLIRBufferizationTransforms
     MLIRGPUDialect
@@ -25,6 +26,7 @@ iree_cc_library(
     MLIRMemRefDialect
     MLIRSupport
     iree::compiler::Codegen::Dialect::GPU::IR::IREEGPUDialect
+    iree::compiler::Dialect::HAL::IR
   PUBLIC
 )
 
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
index 80c668c5c458..db6f1b762d0c 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
@@ -184,6 +184,7 @@ static ReorderWorkgroupsStrategy getReorderWorkgroupsStrategy(
 
 static void addBufferizePasses(OpPassManager &funcPassManager) {
   funcPassManager.addPass(createROCDLConfigureBufferInstructionsPass());
+  funcPassManager.addPass(createGPUBubbleResourceCastsPass());
   BufferizationOptions::AllocationFn allocationFn = gpuAllocationFn;
   BufferizationOptions::MemCpyFn memcpyFn = gpuCopyFn;
   addIREEComprehensiveBufferizePasses(funcPassManager, allocationFn, memcpyFn);
@@ -334,6 +335,7 @@ static void addGPUBufferizePasses(OpPassManager &funcPassManager) {
   funcPassManager.addPass(bufferization::createEmptyTensorToAllocTensorPass());
   funcPassManager.addPass(createGPUInferMemorySpacePass());
   funcPassManager.addPass(createROCDLConfigureBufferInstructionsPass());
+  funcPassManager.addPass(createGPUBubbleResourceCastsPass());
   BufferizationOptions::AllocationFn allocationFn =
       gpuRequireMemSpaceAllocationFn;
   BufferizationOptions::MemCpyFn memcpyFn = [](OpBuilder &builder, Location loc,
@@ -462,6 +464,10 @@ void addGPUTileAndFusePassPipeline(OpPassManager &funcPassManager,
   // hoisting and fusion pass, as well as a lack of a fallback distribution
   // pass.
   funcPassManager.addPass(createIREELoopInvariantCodeMotionPass());
+
+  // Drop resource casts if needed. This is the last possible place to do so
+  // before greedy fusion.
+  funcPassManager.addPass(createGPUBubbleResourceCastsPass());
   {
     OptimizeTensorInsertExtractSlicesPassOptions options;
     options.foldIdentitySlices = true;


Reviews:
user: login: krzysz00
body: I'm missing something: could you explain why we need this and also why it isn't folded into the "can we use buffers for this" checking pass?
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T20:28:56Z
user: login: krzysz00
body: Overall this looks reasonable - I'm much happier with the design here


state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T23:31:09Z
user: login: qedawkins
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T23:39:19Z
user: login: qedawkins
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T19:27:41Z
user: login: krzysz00
body: LGTM - is there a followup patch where we start using this more?
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T21:35:02Z
user: login: qedawkins
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T22:18:45Z

Review comments:
diff_hunk: @@ -0,0 +1,199 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Codegen/Common/GPU/Passes.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
+#include "iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler {
+
+#define GEN_PASS_DEF_GPUBUBBLERESOURCECASTSPASS
+#include "iree/compiler/Codegen/Common/GPU/Passes.h.inc"
+
+namespace {
+struct GPUBubbleResourceCastsPass final
+    : impl::GPUBubbleResourceCastsPassBase<GPUBubbleResourceCastsPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+/// Walk the producers of |input| and return the hal.interface.binding_subspan
+/// that produces it if the binding is readonly.
+static Operation *isTriviallyProducedByReadOnlyViewLike(Value input) {
+  Value next = input;
+  while (auto slice = next.getDefiningOp<tensor::ExtractSliceOp>()) {
path: compiler/src/iree/compiler/Codegen/Common/GPU/GPUBubbleResourceCasts.cpp
user: login: krzysz00
body: Are there other view-like ops this could be? reshape and friends?
created_at: 2025-04-28T23:24:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,199 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Codegen/Common/GPU/Passes.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
+#include "iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler {
+
+#define GEN_PASS_DEF_GPUBUBBLERESOURCECASTSPASS
+#include "iree/compiler/Codegen/Common/GPU/Passes.h.inc"
+
+namespace {
+struct GPUBubbleResourceCastsPass final
+    : impl::GPUBubbleResourceCastsPassBase<GPUBubbleResourceCastsPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+/// Walk the producers of |input| and return the hal.interface.binding_subspan
+/// that produces it if the binding is readonly.
+static Operation *isTriviallyProducedByReadOnlyViewLike(Value input) {
+  Value next = input;
+  while (auto slice = next.getDefiningOp<tensor::ExtractSliceOp>()) {
+    next = slice.getSource();
+  }
+
+  auto tensorLoad = next.getDefiningOp<IREE::TensorExt::DispatchTensorLoadOp>();
+  if (!tensorLoad) {
+    return nullptr;
+  }
+
+  auto binding = tensorLoad.getSource()
+                     .getDefiningOp<IREE::HAL::InterfaceBindingSubspanOp>();
+  if (!binding ||
+      cast<IREE::TensorExt::DispatchTensorType>(binding.getType())
+              .getAccess() != IREE::TensorExt::TensorAccess::ReadOnly) {
+    return nullptr;
+  }
+  return binding;
+}
+
+namespace {
+struct BubbleResourceCastPattern
+    : public OpRewritePattern<IREE::GPU::BufferResourceCastOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(IREE::GPU::BufferResourceCastOp castOp,
+                                PatternRewriter &rewriter) const override {
+    // Skip ops with cache swizzle strides because the stride could be
+    // rendered invalid by the swap.
+    // TODO: Support bubbling swizzled casts in cases where it makes sense.
+    if (castOp.getCacheSwizzleStride()) {
path: compiler/src/iree/compiler/Codegen/Common/GPU/GPUBubbleResourceCasts.cpp
user: login: krzysz00
body: Do we have examples here?
created_at: 2025-04-28T23:27:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,199 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Codegen/Common/GPU/Passes.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
+#include "iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler {
+
+#define GEN_PASS_DEF_GPUBUBBLERESOURCECASTSPASS
+#include "iree/compiler/Codegen/Common/GPU/Passes.h.inc"
+
+namespace {
+struct GPUBubbleResourceCastsPass final
+    : impl::GPUBubbleResourceCastsPassBase<GPUBubbleResourceCastsPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+/// Walk the producers of |input| and return the hal.interface.binding_subspan
+/// that produces it if the binding is readonly.
+static Operation *isTriviallyProducedByReadOnlyViewLike(Value input) {
+  Value next = input;
+  while (auto slice = next.getDefiningOp<tensor::ExtractSliceOp>()) {
path: compiler/src/iree/compiler/Codegen/Common/GPU/GPUBubbleResourceCasts.cpp
user: login: qedawkins
body: This is a big type switch on the ops I could think of for now, pad and reshape included down below. A notable omission is cast because that's better handled as a canonicalization due to the way we want static info to propagate across the IR.

A whitelist made more sense to me to start but if there is a sensible general alternative we can switch in the future.
created_at: 2025-04-28T23:39:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,199 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Codegen/Common/GPU/Passes.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
+#include "iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler {
+
+#define GEN_PASS_DEF_GPUBUBBLERESOURCECASTSPASS
+#include "iree/compiler/Codegen/Common/GPU/Passes.h.inc"
+
+namespace {
+struct GPUBubbleResourceCastsPass final
+    : impl::GPUBubbleResourceCastsPassBase<GPUBubbleResourceCastsPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+/// Walk the producers of |input| and return the hal.interface.binding_subspan
+/// that produces it if the binding is readonly.
+static Operation *isTriviallyProducedByReadOnlyViewLike(Value input) {
+  Value next = input;
+  while (auto slice = next.getDefiningOp<tensor::ExtractSliceOp>()) {
+    next = slice.getSource();
+  }
+
+  auto tensorLoad = next.getDefiningOp<IREE::TensorExt::DispatchTensorLoadOp>();
+  if (!tensorLoad) {
+    return nullptr;
+  }
+
+  auto binding = tensorLoad.getSource()
+                     .getDefiningOp<IREE::HAL::InterfaceBindingSubspanOp>();
+  if (!binding ||
+      cast<IREE::TensorExt::DispatchTensorType>(binding.getType())
+              .getAccess() != IREE::TensorExt::TensorAccess::ReadOnly) {
+    return nullptr;
+  }
+  return binding;
+}
+
+namespace {
+struct BubbleResourceCastPattern
+    : public OpRewritePattern<IREE::GPU::BufferResourceCastOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(IREE::GPU::BufferResourceCastOp castOp,
+                                PatternRewriter &rewriter) const override {
+    // Skip ops with cache swizzle strides because the stride could be
+    // rendered invalid by the swap.
+    // TODO: Support bubbling swizzled casts in cases where it makes sense.
+    if (castOp.getCacheSwizzleStride()) {
path: compiler/src/iree/compiler/Codegen/Common/GPU/GPUBubbleResourceCasts.cpp
user: login: qedawkins
body: I added a short example to the comment
created_at: 2025-04-29T19:27:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -456,4 +456,44 @@ def IREEGPU_YieldOp : Op<IREEGPU_Dialect, "yield", [
       [{  attr-dict ($values^ `:` type($values))? }];
 }
 
+//===----------------------------------------------------------------------===//
+//
+// AMD Specific Operations
+//
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// BufferResourceCastOp
+//===----------------------------------------------------------------------===//
+
+def IREEGPU_BufferResourceCastOp : Op<IREEGPU_Dialect, "buffer_resource_cast", [
+  Pure,
+  AllTypesMatch<["input", "result"]>]> {
+  let summary = "Represents a cast to addr_space<7> (buffer resource) before bufferization";
+  let description = [{
+    Nominal cast of a tensor to AMDGPU buffer resource memory space before
+    bufferization. This op takes the parameters with which to perform the cast
+    if |input| bufferizes to `storage_buffer` memory space. If |input| resolves
+    to any other memory space this op is silently dropped and has no effect.
+
+    If |cache_swizzle_stride| is present, there is verification before
path: compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.td
user: login: krzysz00
body: Why're you using `|param|`? Is that a style thing I'm unaware of?
created_at: 2025-04-29T21:32:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -456,4 +456,44 @@ def IREEGPU_YieldOp : Op<IREEGPU_Dialect, "yield", [
       [{  attr-dict ($values^ `:` type($values))? }];
 }
 
+//===----------------------------------------------------------------------===//
+//
+// AMD Specific Operations
+//
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// BufferResourceCastOp
+//===----------------------------------------------------------------------===//
+
+def IREEGPU_BufferResourceCastOp : Op<IREEGPU_Dialect, "buffer_resource_cast", [
+  Pure,
+  AllTypesMatch<["input", "result"]>]> {
+  let summary = "Represents a cast to addr_space<7> (buffer resource) before bufferization";
+  let description = [{
+    Nominal cast of a tensor to AMDGPU buffer resource memory space before
+    bufferization. This op takes the parameters with which to perform the cast
+    if |input| bufferizes to `storage_buffer` memory space. If |input| resolves
+    to any other memory space this op is silently dropped and has no effect.
+
+    If |cache_swizzle_stride| is present, there is verification before
path: compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.td
user: login: qedawkins
body: It's pretty common through the rest of the codebase, e.g.

- https://github.com/iree-org/iree/blob/72f11570a2582f2e1703dc00eda6780fcb8f8c5d/runtime/src/iree/vm/buffer.h#L76-L82
- https://github.com/iree-org/iree/blob/72f11570a2582f2e1703dc00eda6780fcb8f8c5d/compiler/src/iree/compiler/Dialect/Util/IR/UtilOps.td#L345

but isn't formal style as far as I am aware, I just like it over backticks or quotes.
created_at: 2025-04-29T22:18:45Z
author_association: CONTRIBUTOR

PR comments:
user: login: qedawkins
created_at: 2025-04-23T23:53:29Z
body: Four main reasons for the op:

1. We need a way to represent casts other than on the binding. For example, we can do the cast inside a hot loop to potentially reduce VGPR pressure and turn VALU operations into SALU ops.
2. Setting the value to use for cache swizzle is easier to do in tensor land before we've tiled and distributed than trying to reverse engineer it during bufferization.
3. The op is marginally more future proof against future patterns on the HAL/Stream side that could collapse a binding into a single one with multiple offsets. We'd still need some kind of "no-alias" equivalent info though.
4. Admittedly specific to my use case, but I want to be able to set a cache swizzle value from hand-written IR that propagates up to the buffer resource cast. The only way I could think to do that was with an op that represents the cast but folds away if there are any unexpected producers.

I might be able to think of more reasons, but I'm quite sure we want the op based on what I've seen Triton do with buffer instructions.

> why it isn't folded into the "can we use buffers for this" checking pass?

Good q, I wasn't quite sure how best to do that yet, and because I was hand coding the op in my use case I kind of cheated with a separate pass. Rolling the DropResourceCasts pass into the checking pass makes sense after I just reviewed the code there.
author_association: CONTRIBUTOR
user: login: qedawkins
created_at: 2025-04-28T23:19:35Z
body: @krzysz00 still missing some more tests, but this should be what we discussed. Lot's of room for improvement in the future too, but PTAL whenever you have a moment.
author_association: CONTRIBUTOR
user: login: qedawkins
created_at: 2025-04-29T22:22:04Z
body: > is there a followup patch where we start using this more?

Implemented? No :P

Yes in the semi near future though. I need to work out/benchmark how to pick the right cache swizzle value before we get benefit of on-by-default, but the hope is it should be easy to generate after operand promotion (maybe earlier before workgroup tiling).
author_association: CONTRIBUTOR
user: login: qedawkins
created_at: 2025-04-30T16:42:00Z
body: All non-bazel tests passed last night. Going to merge through remaining ONNX tests.
author_association: CONTRIBUTOR