refactor: Allow accessing the underlying buffer for TensorMap (#940)

* Make `as_tensor_map` optional

* Revert debug setting

* Add 1D TMA load

* Remove leftover print

* Add 1D TMA store

Author: wingertge

crates/cubecl-attention/src/components/args.rs

@@ -1,6 +1,9 @@
 use cubecl::prelude::*;
 use cubecl_core::{self as cubecl};
-use cubecl_std::tensor::r#virtual::{VirtualTensorOperations, VirtualTensorOperationsExpand};
+use cubecl_std::{
+    CubeOption, CubeOptionExpand,
+    tensor::r#virtual::{VirtualTensorOperations, VirtualTensorOperationsExpand},
+};
 
 use crate::components::{
     line_size::AttentionLineSizes, problem::AttentionProblem, selection::AttentionSelection,
@@ -88,17 +91,17 @@ pub trait AttentionArgs: Send + Sync + 'static + Clone {
     /// Reinterpret query as tensor map
     fn as_tensor_map_query<Q: Float, K: Float, V: Float, O: Float>(
         state: &Self::State<Q, K, V, O>,
-    ) -> TensorMap<Q>;
+    ) -> CubeOption<TensorMap<Q>>;
 
     /// Reinterpret key as tensor map
     fn as_tensor_map_key<Q: Float, K: Float, V: Float, O: Float>(
         state: &Self::State<Q, K, V, O>,
-    ) -> TensorMap<K>;
+    ) -> CubeOption<TensorMap<K>>;
 
     /// Reinterpret value as tensor map
     fn as_tensor_map_value<Q: Float, K: Float, V: Float, O: Float>(
         state: &Self::State<Q, K, V, O>,
-    ) -> TensorMap<V>;
+    ) -> CubeOption<TensorMap<V>>;
 
     /// Write the line to the output at the given coordinate using the state.
     fn write_out<Q: Float, K: Float, V: Float, O: Float>(
@@ -290,11 +293,8 @@ impl<Q: Float, K: Float, V: Float, O: Float, MA: AttentionArgs> VirtualTensorOpe
         TensorOutputExpand::__expand_buffer_len_method(self.clone(), scope)
     }
 
-    fn __expand_as_tensor_map_method(
-        &self,
-        _scope: &mut Scope,
-    ) -> ExpandElementTyped<TensorMap<O>> {
-        unimplemented!("TensorOutputExpand can't be turned into a tensor map");
+    fn __expand_as_tensor_map_method(&self, scope: &mut Scope) -> CubeOptionExpand<TensorMap<O>> {
+        CubeOption::__expand_new_None(scope)
     }
 }
 
@@ -367,7 +367,7 @@ impl<Q: Float, K: Float, V: Float, O: Float, MA: AttentionArgs> VirtualTensorOpe
         TensorQueryExpand::__expand_buffer_len_method(self.clone(), scope)
     }
 
-    fn __expand_as_tensor_map_method(&self, scope: &mut Scope) -> ExpandElementTyped<TensorMap<Q>> {
+    fn __expand_as_tensor_map_method(&self, scope: &mut Scope) -> CubeOptionExpand<TensorMap<Q>> {
         TensorQueryExpand::__expand_as_tensor_map_method(self.clone(), scope)
     }
 }
@@ -441,7 +441,7 @@ impl<Q: Float, K: Float, V: Float, O: Float, MA: AttentionArgs> VirtualTensorOpe
         TensorKeyExpand::__expand_buffer_len_method(self.clone(), scope)
     }
 
-    fn __expand_as_tensor_map_method(&self, scope: &mut Scope) -> ExpandElementTyped<TensorMap<K>> {
+    fn __expand_as_tensor_map_method(&self, scope: &mut Scope) -> CubeOptionExpand<TensorMap<K>> {
         TensorKeyExpand::__expand_as_tensor_map_method(self.clone(), scope)
     }
 }
@@ -515,7 +515,7 @@ impl<Q: Float, K: Float, V: Float, O: Float, MA: AttentionArgs> VirtualTensorOpe
         TensorValueExpand::__expand_buffer_len_method(self.clone(), scope)
     }
 
-    fn __expand_as_tensor_map_method(&self, scope: &mut Scope) -> ExpandElementTyped<TensorMap<V>> {
+    fn __expand_as_tensor_map_method(&self, scope: &mut Scope) -> CubeOptionExpand<TensorMap<V>> {
         TensorValueExpand::__expand_as_tensor_map_method(self.clone(), scope)
     }
 }
@@ -606,7 +606,7 @@ impl<Q: Float, K: Float, V: Float, O: Float, MA: AttentionArgs> TensorQuery<Q, K
     }
 
     /// Get the buffer length of the tensor.
-    pub fn as_tensor_map(&self) -> TensorMap<Q> {
+    pub fn as_tensor_map(&self) -> CubeOption<TensorMap<Q>> {
         unsafe { MA::as_tensor_map_query(&(*self.state)) }
     }
 
@@ -660,7 +660,7 @@ impl<Q: Float, K: Float, V: Float, O: Float, MA: AttentionArgs> TensorKey<Q, K,
     }
 
     /// Get the buffer length of the tensor.
-    pub fn as_tensor_map(&self) -> TensorMap<K> {
+    pub fn as_tensor_map(&self) -> CubeOption<TensorMap<K>> {
         unsafe { MA::as_tensor_map_key(&(*self.state)) }
     }
 
@@ -714,7 +714,7 @@ impl<Q: Float, K: Float, V: Float, O: Float, MA: AttentionArgs> TensorValue<Q, K
     }
 
     /// Get the buffer length of the tensor.
-    pub fn as_tensor_map(&self) -> TensorMap<V> {
+    pub fn as_tensor_map(&self) -> CubeOption<TensorMap<V>> {
         unsafe { MA::as_tensor_map_value(&(*self.state)) }
     }
 
@@ -885,26 +885,20 @@ impl AttentionArgs for TensorArgs {
 
     fn as_tensor_map_query<Q: Float, K: Float, V: Float, O: Float>(
         _state: &Self::State<Q, K, V, O>,
-    ) -> TensorMap<Q> {
-        comptime!(unimplemented!("Can't use `TensorArgs` as `TensorMap`"));
-        #[allow(unreachable_code)]
-        TensorMap::dummy()
+    ) -> CubeOption<TensorMap<Q>> {
+        CubeOption::new_None()
     }
 
     fn as_tensor_map_key<Q: Float, K: Float, V: Float, O: Float>(
         _state: &Self::State<Q, K, V, O>,
-    ) -> TensorMap<K> {
-        comptime!(unimplemented!("Can't use `TensorArgs` as `TensorMap`"));
-        #[allow(unreachable_code)]
-        TensorMap::dummy()
+    ) -> CubeOption<TensorMap<K>> {
+        CubeOption::new_None()
     }
 
     fn as_tensor_map_value<Q: Float, K: Float, V: Float, O: Float>(
         _state: &Self::State<Q, K, V, O>,
-    ) -> TensorMap<V> {
-        comptime!(unimplemented!("Can't use `TensorArgs` as `TensorMap`"));
-        #[allow(unreachable_code)]
-        TensorMap::dummy()
+    ) -> CubeOption<TensorMap<V>> {
+        CubeOption::new_None()
     }
 
     fn shape_query<Q: Float, K: Float, V: Float, O: Float>(

crates/cubecl-convolution/src/components/global/memory/tma.rs

@@ -22,7 +22,7 @@ impl<E: Numeric> Im2colTmaReader<E> {
         spatial_offsets: Sequence<u32>,
         k_offset: u32,
     ) -> Im2colTmaReader<E> {
-        let map = tensor.as_tensor_map();
+        let map = tensor.as_tensor_map().unwrap();
 
         Im2colTmaReader::<E> {
             tensor: map,

crates/cubecl-convolution/src/components/global/multi_stage/tma/convolution.rs

@@ -222,7 +222,7 @@ where
     ) -> Self::RhsGlobalReader {
         let (x_offset, y_offset) = offset;
         Self::RhsGlobalReader::new(
-            rhs.as_tensor_map(),
+            rhs.as_tensor_map().unwrap(),
             x_offset,
             y_offset,
             runtime_args,

crates/cubecl-convolution/src/components/global/single_stage/tma/convolution.rs

@@ -151,7 +151,7 @@ where
     ) -> Self::RhsGlobalReader {
         let (x_offset, y_offset) = offset;
         Self::RhsGlobalReader::new(
-            rhs.as_tensor_map(),
+            rhs.as_tensor_map().unwrap(),
             x_offset,
             y_offset,
             runtime_args,

crates/cubecl-core/src/codegen/integrator.rs

@@ -14,15 +14,15 @@ pub struct KernelIntegrator {
     expansion: KernelExpansion,
     buffer_bindings: Vec<Binding>,
     scalar_bindings: Vec<ScalarBinding>,
-    tensor_maps: Vec<Id>,
+    tensor_maps: Vec<Binding>,
 }
 
 /// The information necessary to compile a [kernel definition](KernelDefinition).
 #[derive(Clone)]
 pub struct KernelExpansion {
     pub buffers: Vec<BufferInfo>,
     pub scalars: Vec<ScalarInfo>,
-    pub tensor_maps: Vec<Id>,
+    pub tensor_maps: Vec<BufferInfo>,
     pub scope: Scope,
 }
 
@@ -143,8 +143,15 @@ impl KernelIntegrator {
     }
 
     fn register_tensor_maps(&mut self) {
-        for id in self.expansion.tensor_maps.drain(..) {
-            self.tensor_maps.push(id);
+        for buffer in self.expansion.tensor_maps.drain(..) {
+            self.tensor_maps.push(Binding {
+                id: buffer.id,
+                ty: buffer.item,
+                visibility: buffer.visibility,
+                location: Location::Storage,
+                has_extended_meta: buffer.has_extended_meta,
+                size: None,
+            });
         }
     }
 }

crates/cubecl-core/src/compute/builder.rs

@@ -5,9 +5,7 @@ use std::{
 
 use alloc::collections::BTreeMap;
 
-use cubecl_ir::{
-    ExpandElement, Scope, SemanticType, StorageType, TargetProperties, Variable, VariableKind,
-};
+use cubecl_ir::{ExpandElement, Scope, StorageType, TargetProperties, Variable, VariableKind};
 use cubecl_runtime::config::{GlobalConfig, compilation::CompilationLogLevel};
 
 use crate::ir::{Id, Type};
@@ -23,7 +21,7 @@ pub struct KernelBuilder {
     pub scope: Scope,
     buffers: Vec<BufferInfo>,
     scalars: BTreeMap<StorageType, usize>,
-    tensor_maps: Vec<Id>,
+    tensor_maps: Vec<BufferInfo>,
 }
 
 static DEBUG: AtomicI8 = AtomicI8::new(-1);
@@ -54,13 +52,27 @@ impl KernelBuilder {
     }
 
     /// Register a tensor map and return the [element](ExpandElement) to be used for kernel expansion.
-    pub fn tensor_map(&mut self) -> ExpandElement {
+    pub fn input_tensor_map(&mut self, item: Type) -> ExpandElement {
         let id = self.buffer_id();
-        self.tensor_maps.push(id);
-        ExpandElement::Plain(Variable::new(
-            VariableKind::TensorMap(id),
-            Type::semantic(SemanticType::TensorMap),
-        ))
+        self.tensor_maps.push(BufferInfo {
+            id,
+            item,
+            visibility: Visibility::ReadWrite,
+            has_extended_meta: true,
+        });
+        ExpandElement::Plain(Variable::new(VariableKind::TensorMapInput(id), item))
+    }
+
+    /// Register a tensor map and return the [element](ExpandElement) to be used for kernel expansion.
+    pub fn output_tensor_map(&mut self, item: Type) -> ExpandElement {
+        let id = self.buffer_id();
+        self.tensor_maps.push(BufferInfo {
+            id,
+            item,
+            visibility: Visibility::Read,
+            has_extended_meta: true,
+        });
+        ExpandElement::Plain(Variable::new(VariableKind::TensorMapOutput(id), item))
     }
 
     /// Register an input array and return the [element](ExpandElement) to be used for kernel expansion.

crates/cubecl-core/src/compute/kernel.rs

@@ -156,7 +156,7 @@ source:
 #[allow(missing_docs)]
 pub struct KernelDefinition {
     pub buffers: Vec<Binding>,
-    pub tensor_maps: Vec<Id>,
+    pub tensor_maps: Vec<Binding>,
     pub scalars: Vec<ScalarBinding>,
     pub cube_dim: CubeDim,
     pub body: Scope,

crates/cubecl-core/src/frontend/barrier.rs

@@ -240,6 +240,7 @@ macro_rules! tensor_map_load_im2col {
     };
 }
 
+tensor_map_load!(1, x);
 tensor_map_load!(2, y, x);
 tensor_map_load!(3, z, y, x);
 tensor_map_load!(4, w, z, y, x);