[Codegen] Add bufferization support for CastToRaggedShapeOp and LinearizeRaggedDimsOp (#24193) This is PR 4/N of landing block sparse changes that are captured in this branch : https://github.com/MaheshRavishankar/iree/tree/users/MaheshRavishankar/blockSparseCommits Adds `BufferizableOpInterfac`e external models for two tensor-ext ops that describe ragged shapes: - `iree_tensor_ext.cast_to_ragged_shape` — reinterprets a source buffer as a rank-(n+1) ragged buffer by attaching a RaggedShapeAttr as the memref layout. -`iree_tensor_ext.linearize_ragged_dims` — the inverse view: collapses the ragged dimensions of a source buffer back into a single dense dimension and drops the ragged layout. Both ops are pure views, so the interface implementations: - Report the operands as reads and never as writes (including `resultBufferizesToMemoryWrite = false`). - Set `mustBufferizeInPlace = true` and mark the result as `BufferRelation::Equivalent` to source, so bufferization does not introduce a copy. - For `cast_to_ragged_shape`, the result memref's layout is the ragged encoding promoted from the tensor encoding; operand buffers use the identity layout. - For `linearize_ragged_dims`, both operand and result memrefs use the identity layout; the result simply drops the ragged encoding. `bufferize` replaces the tensor-typed op with an identical memref-typed version of itself — no data is moved.
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 e09f4bd..314bcb2 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
@@ -3154,3 +3154,227 @@ // CHECK-LABEL: func.func @constant_to_buffer_rocm // CHECK: %[[CST:.+]] = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32> // CHECK: bufferization.to_buffer %[[CST]] {{.*}} : tensor<4xi32> to memref<4xi32, #gpu.address_space<constant>> + +// ----- + +#pipeline_layout = #hal.pipeline.layout<bindings = [ + #hal.pipeline.binding<storage_buffer>, + #hal.pipeline.binding<storage_buffer> +]> + +func.func @cast_to_ragged_shape_static() -> tensor<10x3x?x30xf32, #iree_tensor_ext.ragged_shape<1>> { + %c0 = arith.constant 0 : index + %source_binding = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<10x20x30xf32>> + %column_lengths_binding = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<4xindex>> + %source = iree_tensor_ext.dispatch.tensor.load %source_binding, offsets=[0, 0, 0], sizes=[10, 20, 30], strides=[1, 1, 1] + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<10x20x30xf32>> -> tensor<10x20x30xf32> + %column_lengths = iree_tensor_ext.dispatch.tensor.load %column_lengths_binding, offsets=[0], sizes=[4], strides=[1] + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<4xindex>> -> tensor<4xindex> + %0 = iree_tensor_ext.cast_to_ragged_shape %source ragged_dim(1) column_lengths(%column_lengths) + : (tensor<10x20x30xf32>, tensor<4xindex>) + -> tensor<10x3x?x30xf32, #iree_tensor_ext.ragged_shape<1>> + return %0 : tensor<10x3x?x30xf32, #iree_tensor_ext.ragged_shape<1>> +} + +// CHECK-LABEL: func.func @cast_to_ragged_shape_static() +// CHECK: %[[SOURCE:.+]] = hal.interface.binding.subspan {{.+}} binding(0) +// CHECK-SAME: : memref<10x20x30xf32, #hal.descriptor_type<storage_buffer>> +// CHECK: %[[COLUMN_LENGTHS:.+]] = hal.interface.binding.subspan {{.+}} binding(1) +// CHECK-SAME: : memref<4xindex, #hal.descriptor_type<storage_buffer>> +// CHECK: %[[RAGGED:.+]] = iree_tensor_ext.cast_to_ragged_shape %[[SOURCE]] +// CHECK-SAME: ragged_dim(1) column_lengths(%[[COLUMN_LENGTHS]]) +// CHECK-SAME: -> memref<10x3x?x30xf32, #iree_tensor_ext.ragged_shape<1>, #hal.descriptor_type<storage_buffer>> +// CHECK: %[[RAGGED_TENSOR:.+]] = bufferization.to_tensor %[[RAGGED]] +// CHECK: return %[[RAGGED_TENSOR]] + +// ----- + +#pipeline_layout = #hal.pipeline.layout<bindings = [ + #hal.pipeline.binding<storage_buffer>, + #hal.pipeline.binding<storage_buffer> +]> +func.func @cast_to_ragged_shape_dynamic(%num_ragged_rows : index, %d0 : index, %d1 : index, %d2 : index) + -> tensor<?x?x?x?xf32, #iree_tensor_ext.ragged_shape<1>> { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %source_binding = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?x?xf32>>{%d0, %d1, %d2} + %num_ragged_rows_p1 = arith.addi %num_ragged_rows, %c1 : index + %column_lengths_binding = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?xindex>>{%num_ragged_rows_p1} + %source = iree_tensor_ext.dispatch.tensor.load %source_binding, offsets=[0, 0, 0], sizes=[%d0, %d1, %d2], strides=[1, 1, 1] + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?x?xf32>>{%d0, %d1, %d2} -> tensor<?x?x?xf32> + %column_lengths = iree_tensor_ext.dispatch.tensor.load %column_lengths_binding, offsets=[0], sizes=[%num_ragged_rows_p1], strides=[1] + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?xindex>>{%num_ragged_rows_p1} -> tensor<?xindex> + %0 = iree_tensor_ext.cast_to_ragged_shape %source ragged_dim(1) column_lengths(%column_lengths) + num_ragged_rows(%num_ragged_rows) + : (tensor<?x?x?xf32>{%d0, %d1, %d2}, tensor<?xindex>) + -> tensor<?x?x?x?xf32, #iree_tensor_ext.ragged_shape<1>> + return %0 : tensor<?x?x?x?xf32, #iree_tensor_ext.ragged_shape<1>> +} + +// CHECK-LABEL: func.func @cast_to_ragged_shape_dynamic( +// CHECK-SAME: %[[NUM_RAGGED_ROWS:[A-Za-z0-9]+]]: index +// CHECK-SAME: %[[D0:[A-Za-z0-9]+]]: index +// CHECK-SAME: %[[D1:[A-Za-z0-9]+]]: index +// CHECK-SAME: %[[D2:[A-Za-z0-9]+]]: index +// CHECK: %[[C1:.+]] = arith.constant 1 : index +// CHECK: %[[SOURCE:.+]] = hal.interface.binding.subspan {{.+}} binding(0) +// CHECK-SAME: : memref<?x?x?xf32, #hal.descriptor_type<storage_buffer>>{%[[D0]], %[[D1]], %[[D2]]} +// CHECK: %[[NUM_RAGGED_ROWS_P1:.+]] = arith.addi %[[NUM_RAGGED_ROWS]], %[[C1]] : index +// CHECK: %[[COLUMN_LENGTHS:.+]] = hal.interface.binding.subspan {{.+}} binding(1) +// CHECK-SAME: : memref<?xindex, #hal.descriptor_type<storage_buffer>> +// CHECK: %[[RAGGED:.+]] = iree_tensor_ext.cast_to_ragged_shape %[[SOURCE]] +// CHECK-SAME: ragged_dim(1) column_lengths(%[[COLUMN_LENGTHS]]) +// CHECK-SAME: -> memref<?x?x?x?xf32, #iree_tensor_ext.ragged_shape<1>, #hal.descriptor_type<storage_buffer>> +// CHECK: %[[RAGGED_TENSOR:.+]] = bufferization.to_tensor %[[RAGGED]] +// CHECK: return %[[RAGGED_TENSOR]] +// ----- + +#pipeline_layout = #hal.pipeline.layout<bindings = [ + #hal.pipeline.binding<storage_buffer>, + #hal.pipeline.binding<storage_buffer> +]> +func.func @cast_to_ragged_shape_dynamic_dims_known_num_rows( + %d0 : index, %d1 : index, %d2 : index) -> tensor<?x3x?x?xf32, #iree_tensor_ext.ragged_shape<1>> { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %source_binding = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?x?xf32>>{%d0, %d1, %d2} + %column_lengths_binding = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<4xindex>> + %source = iree_tensor_ext.dispatch.tensor.load %source_binding, offsets=[0, 0, 0], sizes=[%d0, %d1, %d2], strides=[1, 1, 1] + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?x?xf32>>{%d0, %d1, %d2} -> tensor<?x?x?xf32> + %column_lengths = iree_tensor_ext.dispatch.tensor.load %column_lengths_binding, offsets=[0], sizes=[4], strides=[1] + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<4xindex>> -> tensor<4xindex> + %0 = iree_tensor_ext.cast_to_ragged_shape %source ragged_dim(1) column_lengths(%column_lengths) + : (tensor<?x?x?xf32>{%d0, %d1, %d2}, tensor<4xindex>) + -> tensor<?x3x?x?xf32, #iree_tensor_ext.ragged_shape<1>> + return %0 : tensor<?x3x?x?xf32, #iree_tensor_ext.ragged_shape<1>> +} + +// CHECK-LABEL: func.func @cast_to_ragged_shape_dynamic_dims_known_num_rows( +// CHECK-SAME: %[[D0:[A-Za-z0-9]+]]: index +// CHECK-SAME: %[[D1:[A-Za-z0-9]+]]: index +// CHECK-SAME: %[[D2:[A-Za-z0-9]+]]: index +// CHECK: %[[SOURCE:.+]] = hal.interface.binding.subspan {{.+}} binding(0) +// CHECK-SAME: : memref<?x?x?xf32, #hal.descriptor_type<storage_buffer>>{%[[D0]], %[[D1]], %[[D2]]} +// CHECK-DAG: %[[COLUMN_LENGTHS:.+]] = hal.interface.binding.subspan {{.+}} binding(1) +// CHECK-SAME: : memref<4xindex, #hal.descriptor_type<storage_buffer>> +// CHECK: %[[RAGGED:.+]] = iree_tensor_ext.cast_to_ragged_shape %[[SOURCE]] ragged_dim(1) column_lengths(%[[COLUMN_LENGTHS]]) +// CHECK-SAME: : (memref<?x?x?xf32, #hal.descriptor_type<storage_buffer>>{%[[D0]], %[[D1]], %[[D2]]}, memref<4xindex, #hal.descriptor_type<storage_buffer>>) +// CHECK-SAME: -> memref<?x3x?x?xf32, #iree_tensor_ext.ragged_shape<1>, #hal.descriptor_type<storage_buffer>> +// CHECK: %[[RAGGED_TENSOR:.+]] = bufferization.to_tensor %[[RAGGED]] +// CHECK: return %[[RAGGED_TENSOR]] + +// ----- + +#pipeline_layout = #hal.pipeline.layout<bindings = [ + #hal.pipeline.binding<storage_buffer>, + #hal.pipeline.binding<storage_buffer> +]> +func.func @linearize_ragged_dims_static() -> tensor<10x20x30xf32> { + %c0 = arith.constant 0 : index + %source_binding = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<10x20x30xf32>> + %column_lengths_binding = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<4xindex>> + %source = iree_tensor_ext.dispatch.tensor.load %source_binding, offsets=[0, 0, 0], sizes=[10, 20, 30], strides=[1, 1, 1] + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<10x20x30xf32>> -> tensor<10x20x30xf32> + %column_lengths = iree_tensor_ext.dispatch.tensor.load %column_lengths_binding, offsets=[0], sizes=[4], strides=[1] + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<4xindex>> -> tensor<4xindex> + %ragged = iree_tensor_ext.cast_to_ragged_shape %source ragged_dim(1) column_lengths(%column_lengths) + : (tensor<10x20x30xf32>, tensor<4xindex>) + -> tensor<10x3x?x30xf32, #iree_tensor_ext.ragged_shape<1>> + %0 = iree_tensor_ext.linearize_ragged_dims %ragged + : tensor<10x3x?x30xf32, #iree_tensor_ext.ragged_shape<1>> -> tensor<10x20x30xf32> + return %0 : tensor<10x20x30xf32> +} + +// CHECK-LABEL: func.func @linearize_ragged_dims_static() +// CHECK: %[[SOURCE:.+]] = hal.interface.binding.subspan {{.+}} binding(0) +// CHECK-SAME: : memref<10x20x30xf32, #hal.descriptor_type<storage_buffer>> +// CHECK: %[[COLUMN_LENGTHS:.+]] = hal.interface.binding.subspan {{.+}} binding(1) +// CHECK-SAME: : memref<4xindex, #hal.descriptor_type<storage_buffer>> +// CHECK: %[[RAGGED:.+]] = iree_tensor_ext.cast_to_ragged_shape %[[SOURCE]] +// CHECK-SAME: ragged_dim(1) column_lengths(%[[COLUMN_LENGTHS]]) +// CHECK-SAME: -> memref<10x3x?x30xf32, #iree_tensor_ext.ragged_shape<1>, #hal.descriptor_type<storage_buffer>> +// CHECK: %[[LINEARIZED:.+]] = iree_tensor_ext.linearize_ragged_dims %[[RAGGED]] +// CHECK-SAME: : memref<10x3x?x30xf32, #iree_tensor_ext.ragged_shape<1>, #hal.descriptor_type<storage_buffer>> -> memref<10x20x30xf32, #hal.descriptor_type<storage_buffer>> +// CHECK: %[[RESULT:.+]] = bufferization.to_tensor %[[LINEARIZED]] +// CHECK: return %[[RESULT]] + +// ----- + +#pipeline_layout = #hal.pipeline.layout<bindings = [ + #hal.pipeline.binding<storage_buffer>, + #hal.pipeline.binding<storage_buffer> +]> +func.func @linearize_ragged_dims_dynamic(%d0 : index, %d1 : index, %d2 : index) -> tensor<?x?x?xf32> { + %c0 = arith.constant 0 : index + %source_binding = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?x?xf32>>{%d0, %d1, %d2} + %column_lengths_binding = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<4xindex>> + %source = iree_tensor_ext.dispatch.tensor.load %source_binding, offsets=[0, 0, 0], sizes=[%d0, %d1, %d2], strides=[1, 1, 1] + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?x?xf32>>{%d0, %d1, %d2} -> tensor<?x?x?xf32> + %column_lengths = iree_tensor_ext.dispatch.tensor.load %column_lengths_binding, offsets=[0], sizes=[4], strides=[1] + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<4xindex>> -> tensor<4xindex> + %ragged = iree_tensor_ext.cast_to_ragged_shape %source ragged_dim(1) column_lengths(%column_lengths) + : (tensor<?x?x?xf32>{%d0, %d1, %d2}, tensor<4xindex>) + -> tensor<?x3x?x?xf32, #iree_tensor_ext.ragged_shape<1>> + %0 = iree_tensor_ext.linearize_ragged_dims %ragged + : tensor<?x3x?x?xf32, #iree_tensor_ext.ragged_shape<1>> -> tensor<?x?x?xf32>{%d0, %d1, %d2} + return %0 : tensor<?x?x?xf32> +} + +// CHECK-LABEL: func.func @linearize_ragged_dims_dynamic( +// CHECK-SAME: %[[D0:[A-Za-z0-9]+]]: index +// CHECK-SAME: %[[D1:[A-Za-z0-9]+]]: index +// CHECK-SAME: %[[D2:[A-Za-z0-9]+]]: index +// CHECK: %[[SOURCE:.+]] = hal.interface.binding.subspan {{.+}} binding(0) +// CHECK-SAME: : memref<?x?x?xf32, #hal.descriptor_type<storage_buffer>>{%[[D0]], %[[D1]], %[[D2]]} +// CHECK: %[[COLUMN_LENGTHS:.+]] = hal.interface.binding.subspan {{.+}} binding(1) +// CHECK: %[[RAGGED:.+]] = iree_tensor_ext.cast_to_ragged_shape %[[SOURCE]] +// CHECK-SAME: ragged_dim(1) column_lengths(%[[COLUMN_LENGTHS]]) +// CHECK-SAME: -> memref<?x3x?x?xf32, #iree_tensor_ext.ragged_shape<1>, #hal.descriptor_type<storage_buffer>> +// CHECK: %[[LINEARIZED:.+]] = iree_tensor_ext.linearize_ragged_dims %[[RAGGED]] +// CHECK-SAME: : memref<?x3x?x?xf32, #iree_tensor_ext.ragged_shape<1>, #hal.descriptor_type<storage_buffer>> -> memref<?x?x?xf32, #hal.descriptor_type<storage_buffer>>{%[[D0]], %[[D1]], %[[D2]]} +// CHECK: %[[RESULT:.+]] = bufferization.to_tensor %[[LINEARIZED]] +// CHECK: return %[[RESULT]] + +// ----- + +// A non-default memory space is carried through both ops. +#pipeline_layout = #hal.pipeline.layout<bindings = [ + #hal.pipeline.binding<storage_buffer> +]> +func.func @cast_to_ragged_shape_workgroup_memory() -> tensor<10x20x30xf32> { + %c0 = arith.constant 0 : index + %column_lengths_binding = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<4xindex>> + %column_lengths = iree_tensor_ext.dispatch.tensor.load %column_lengths_binding, offsets=[0], sizes=[4], strides=[1] + : !iree_tensor_ext.dispatch.tensor<readonly:tensor<4xindex>> -> tensor<4xindex> + %alloc = memref.alloc() : memref<10x20x30xf32, #gpu.address_space<workgroup>> + %source = bufferization.to_tensor %alloc restrict writable + : memref<10x20x30xf32, #gpu.address_space<workgroup>> to tensor<10x20x30xf32> + %ragged = iree_tensor_ext.cast_to_ragged_shape %source ragged_dim(1) column_lengths(%column_lengths) + : (tensor<10x20x30xf32>, tensor<4xindex>) + -> tensor<10x3x?x30xf32, #iree_tensor_ext.ragged_shape<1>> + %0 = iree_tensor_ext.linearize_ragged_dims %ragged + : tensor<10x3x?x30xf32, #iree_tensor_ext.ragged_shape<1>> -> tensor<10x20x30xf32> + return %0 : tensor<10x20x30xf32> +} + +// CHECK-LABEL: func.func @cast_to_ragged_shape_workgroup_memory() +// CHECK: %[[COLUMN_LENGTHS:.+]] = hal.interface.binding.subspan {{.+}} binding(0) +// CHECK: %[[ALLOC:.+]] = memref.alloc() : memref<10x20x30xf32, #gpu.address_space<workgroup>> +// CHECK: %[[RAGGED:.+]] = iree_tensor_ext.cast_to_ragged_shape %[[ALLOC]] +// CHECK-SAME: ragged_dim(1) column_lengths(%[[COLUMN_LENGTHS]]) +// CHECK-SAME: -> memref<10x3x?x30xf32, #iree_tensor_ext.ragged_shape<1>, #gpu.address_space<workgroup>> +// CHECK: %[[LINEARIZED:.+]] = iree_tensor_ext.linearize_ragged_dims %[[RAGGED]] +// CHECK-SAME: -> memref<10x20x30xf32, #gpu.address_space<workgroup>> +// CHECK: %[[RESULT:.+]] = bufferization.to_tensor %[[LINEARIZED]] +// CHECK: return %[[RESULT]]
diff --git a/compiler/src/iree/compiler/Codegen/Interfaces/BufferizationInterfaces.cpp b/compiler/src/iree/compiler/Codegen/Interfaces/BufferizationInterfaces.cpp index 23c0b5c..d8db306 100644 --- a/compiler/src/iree/compiler/Codegen/Interfaces/BufferizationInterfaces.cpp +++ b/compiler/src/iree/compiler/Codegen/Interfaces/BufferizationInterfaces.cpp
@@ -563,6 +563,257 @@ } }; +// Bufferization interface for CastToRaggedShapeOp. The op is a pure view: it +// reinterprets its source as a rank-(n+1) ragged buffer by attaching a +// `RaggedShapeAttr` layout, without copying data. The bufferized op therefore +// aliases its source buffer and must bufferize in place. +struct CastToRaggedShapeBufferizationInterface + : public BufferizableOpInterface::ExternalModel< + CastToRaggedShapeBufferizationInterface, + IREE::TensorExt::CastToRaggedShapeOp> { + + // Both operands are read: `source` for its data and `column_lengths` for + // the per-row sizes of the resulting ragged view. + bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, + const bufferization::AnalysisState &state) const { + return true; + } + + // The op produces a view and never writes to any operand. + bool + bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand, + const bufferization::AnalysisState &state) const { + return false; + } + + // Materializing the result view does not introduce new writes either. + bool resultBufferizesToMemoryWrite( + Operation *op, OpResult opResult, + const bufferization::AnalysisState &state) const { + return false; + } + + // A view must share storage with its source; bufferization cannot introduce + // a copy of the source for this op. `column_lengths` does not alias the + // result, so it has no in-place requirement. + bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand, + const bufferization::AnalysisState &state) const { + auto castToRaggedShapeOp = cast<IREE::TensorExt::CastToRaggedShapeOp>(op); + return &opOperand == &castToRaggedShapeOp.getSourceMutable(); + } + + // The result aliases the `source` operand (same underlying buffer, new + // layout). `column_lengths` is consumed as shape metadata and does not + // alias the result. + bufferization::AliasingValueList + getAliasingValues(Operation *op, OpOperand &opOperand, + const bufferization::AnalysisState &state) const { + auto castToRaggedShapeOp = cast<IREE::TensorExt::CastToRaggedShapeOp>(op); + if (opOperand.get() == castToRaggedShapeOp.getSource()) { + return {{castToRaggedShapeOp.getResult(), BufferRelation::Equivalent}}; + } + return {}; + } + + // Inverse of `getAliasingValues`: the result's only aliasing operand is + // `source`. + bufferization::AliasingOpOperandList + getAliasingOpOperands(Operation *op, Value value, + const bufferization::AnalysisState &state) const { + auto castToRaggedShapeOp = cast<IREE::TensorExt::CastToRaggedShapeOp>(op); + if (value == castToRaggedShapeOp.getResult()) { + return {{&castToRaggedShapeOp.getSourceMutable(), + BufferRelation::Equivalent}}; + } + return {}; + } + + // For the result, the memref type preserves the tensor's ragged encoding by + // promoting it to the memref layout attribute. For `source` and + // `column_lengths` operands, use the default identity layout so the op + // consumes plain buffers. + FailureOr<bufferization::BufferLikeType> + getBufferType(Operation *op, Value value, const BufferizationOptions &options, + const BufferizationState &state, + SmallVector<Value> &invocationStack) const { + auto castToRaggedShapeOp = cast<IREE::TensorExt::CastToRaggedShapeOp>(op); + + auto tensorType = dyn_cast<RankedTensorType>(value.getType()); + if (!tensorType) { + return failure(); + } + + if (value == castToRaggedShapeOp.getResult()) { + // The result is a view of the source, so it keeps the source buffer's + // memory space. + FailureOr<bufferization::BufferLikeType> sourceType = + bufferization::getBufferType(castToRaggedShapeOp.getSource(), options, + state, invocationStack); + if (failed(sourceType)) { + return failure(); + } + return cast<bufferization::BufferLikeType>(MemRefType::get( + tensorType.getShape(), tensorType.getElementType(), + /*layout=*/ + cast<MemRefLayoutAttrInterface>(tensorType.getEncoding()), + cast<BaseMemRefType>(*sourceType).getMemorySpace())); + } + + return cast<bufferization::BufferLikeType>( + MemRefType::get(tensorType.getShape(), tensorType.getElementType())); + } + + // Replace the tensor-typed op with the identical op on memrefs: the source + // and column_lengths operands become buffers and the result takes the + // ragged-layout memref type computed by `getBufferType`. No data is copied. + LogicalResult bufferize(Operation *op, RewriterBase &rewriter, + const bufferization::BufferizationOptions &options, + bufferization::BufferizationState &state) const { + auto castToRaggedShapeOp = cast<IREE::TensorExt::CastToRaggedShapeOp>(op); + + FailureOr<Value> sourceBuffer = + getBuffer(rewriter, castToRaggedShapeOp.getSource(), options, state); + if (failed(sourceBuffer)) { + return failure(); + } + + FailureOr<Value> columnLengths = getBuffer( + rewriter, castToRaggedShapeOp.getColumnLengths(), options, state); + if (failed(columnLengths)) { + return failure(); + } + + auto resultMemRefType = bufferization::getBufferType( + castToRaggedShapeOp.getResult(), options, state); + if (failed(resultMemRefType)) { + return failure(); + } + replaceOpWithNewBufferizedOp<IREE::TensorExt::CastToRaggedShapeOp>( + rewriter, op, resultMemRefType.value(), sourceBuffer.value(), + castToRaggedShapeOp.getRaggedDimAttr(), columnLengths.value(), + castToRaggedShapeOp.getNumRaggedRows(), + castToRaggedShapeOp.getSourceDynamicDims()); + return success(); + } +}; + +// Bufferization interface for LinearizeRaggedDimsOp. The op is the dual view +// of `CastToRaggedShape`: it collapses the ragged dimensions of its source +// back into a single dense dimension and drops the `RaggedShapeAttr` layout. +// No data is moved, so the bufferized op aliases its source and must +// bufferize in place. +struct LinearizeRaggedDimsBufferizationInterface + : public BufferizableOpInterface::ExternalModel< + LinearizeRaggedDimsBufferizationInterface, + IREE::TensorExt::LinearizeRaggedDimsOp> { + + // The source buffer is read to produce the linearized view. + bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, + const bufferization::AnalysisState &state) const { + return true; + } + + // The op is a view and does not write to any operand. + bool + bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand, + const bufferization::AnalysisState &state) const { + return false; + } + + // Materializing the result view does not introduce new writes. + bool resultBufferizesToMemoryWrite( + Operation *op, OpResult opResult, + const bufferization::AnalysisState &state) const { + return false; + } + + // The result must share storage with the source; no copy is permitted. + bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand, + const bufferization::AnalysisState &state) const { + return true; + } + + // The result aliases the `source` operand (same underlying buffer, layout + // without ragged encoding). + bufferization::AliasingValueList + getAliasingValues(Operation *op, OpOperand &opOperand, + const bufferization::AnalysisState &state) const { + auto linearizeOp = cast<IREE::TensorExt::LinearizeRaggedDimsOp>(op); + if (opOperand.get() == linearizeOp.getSource()) { + return {{linearizeOp.getResult(), BufferRelation::Equivalent}}; + } + return {}; + } + + // Inverse of `getAliasingValues`: the result's only aliasing operand is + // `source`. + bufferization::AliasingOpOperandList + getAliasingOpOperands(Operation *op, Value value, + const bufferization::AnalysisState &state) const { + auto linearizeOp = cast<IREE::TensorExt::LinearizeRaggedDimsOp>(op); + if (value == linearizeOp.getResult()) { + return {{&linearizeOp.getSourceMutable(), BufferRelation::Equivalent}}; + } + return {}; + } + + // Both source and result use the default identity layout. The result type + // drops the ragged encoding the source carries as part of collapsing the + // ragged dimensions into one dense dimension. + FailureOr<bufferization::BufferLikeType> + getBufferType(Operation *op, Value value, const BufferizationOptions &options, + const BufferizationState &state, + SmallVector<Value> &invocationStack) const { + auto linearizeOp = cast<IREE::TensorExt::LinearizeRaggedDimsOp>(op); + auto tensorType = dyn_cast<RankedTensorType>(value.getType()); + if (!tensorType) { + return failure(); + } + + // The result is a view of the source, so it keeps the source buffer's + // memory space. + Attribute memorySpace; + if (value == linearizeOp.getResult()) { + FailureOr<bufferization::BufferLikeType> sourceType = + bufferization::getBufferType(linearizeOp.getSource(), options, state, + invocationStack); + if (failed(sourceType)) { + return failure(); + } + memorySpace = cast<BaseMemRefType>(*sourceType).getMemorySpace(); + } + return cast<bufferization::BufferLikeType>( + MemRefType::get(tensorType.getShape(), tensorType.getElementType(), + MemRefLayoutAttrInterface{}, memorySpace)); + } + + // Replace the tensor-typed op with the identical op on memrefs: the source + // memref is reinterpreted as a lower-rank non-ragged memref with the + // provided dynamic result dims. No data is copied. + LogicalResult bufferize(Operation *op, RewriterBase &rewriter, + const bufferization::BufferizationOptions &options, + bufferization::BufferizationState &state) const { + auto linearizeOp = cast<IREE::TensorExt::LinearizeRaggedDimsOp>(op); + + FailureOr<Value> sourceBuffer = + getBuffer(rewriter, linearizeOp.getSource(), options, state); + if (failed(sourceBuffer)) { + return failure(); + } + + auto resultMemRefType = + bufferization::getBufferType(linearizeOp.getResult(), options, state); + if (failed(resultMemRefType)) { + return failure(); + } + + replaceOpWithNewBufferizedOp<IREE::TensorExt::LinearizeRaggedDimsOp>( + rewriter, op, resultMemRefType.value(), sourceBuffer.value(), + linearizeOp.getResultDynamicDims()); + return success(); + } +}; + //===----------------------------------------------------------------------===// // IREE specific post analysis transformations. //===----------------------------------------------------------------------===// @@ -582,6 +833,12 @@ IREE::VectorExt::registerIREEVectorExtBufferizationInterfaces(registry); registry.addExtension( +[](MLIRContext *ctx, IREE::TensorExt::IREETensorExtDialect *dialect) { + // CastToRaggedShapeOp + IREE::TensorExt::CastToRaggedShapeOp::attachInterface< + CastToRaggedShapeBufferizationInterface>(*ctx); + // LinearizeRaggedDimsOp + IREE::TensorExt::LinearizeRaggedDimsOp::attachInterface< + LinearizeRaggedDimsBufferizationInterface>(*ctx); // DispatchTensorLoadOp IREE::TensorExt::DispatchTensorLoadOp::attachInterface< DispatchTensorLoadOpInterface>(*ctx);