Alias allocations with same size elements despite different dtypes

csrc/codegen.cpp

@@ -2762,10 +2762,18 @@ class CudaKernelGenerator : private kir::ConstIrVisitor {
     if (alloc->alias() != nullptr) {
       // Allocate alias another Allocate stmt
       const auto alias_tv = alloc->alias()->buffer()->as<TensorView>();
-      indent() << "// Alias Allocation - " << alloc->memoryType() << "\n";
-      indent() << "auto& " << genVariableName(tv) << " = "
-               << genVariableName(alias_tv) << ";\n";
-
+      if (alias_tv->getDataType() == tv->getDataType()) {
+        indent() << "// Alias Allocation - " << alloc->memoryType() << "\n";
+        indent() << "auto& " << genVariableName(tv) << " = "
+                 << genVariableName(alias_tv) << ";\n";
+      } else {
+        indent() << "// Alias Allocation (changing dtype) - "
+                 << alloc->memoryType() << "\n";
+        indent() << "auto " << genVariableName(tv)
+                 << " = *reinterpret_cast<Array<" << buffer_dtype << ", "
+                 << genInline(size) << ">*>(&" << genVariableName(alias_tv)
+                 << ");\n";
+      }
     } else {
       // Standard Memory Allocation
       switch (tv->getMemoryType()) {

csrc/device_lower/pass/alias_memory.cpp

@@ -1053,7 +1053,18 @@ class ReusableAllocationFinder : private kir::IrVisitor {
         }
 
         // Check if this alloc has the same data type
-        if (alloc_info->data_type != alloc_to_reuse->data_type) {
+        if (alloc_info->mem_type == MemoryType::Local &&
+            isOptionDisabled(DisableOption::ReuseMismatchedTypeRegisters)) {
+          // With this option, registers must have exactly matching dtypes in
+          // order to be re-used
+          if (alloc_info->data_type != alloc_to_reuse->data_type) {
+            continue;
+          }
+        } else if (
+            dataTypeSize(alloc_info->data_type) !=
+            dataTypeSize(alloc_to_reuse->data_type)) {
+          // Behavior for shared or global memory and default behavior for
+          // registers is to re-use if dtypes have same size.
           continue;
         }
 

csrc/options.cpp

@@ -180,7 +180,9 @@ std::unordered_map<DisableOption, std::vector<std::string>> Options<
       {"predicate_elimination", DisableOption::PredicateElimination},
       {"kernel_reuse", DisableOption::KernelReuse},
       {"var_name_remapping", DisableOption::VarNameRemapping},
-      {"welford_vectorization", DisableOption::WelfordVectorization}};
+      {"welford_vectorization", DisableOption::WelfordVectorization},
+      {"reuse_mismatched_type_registers",
+       DisableOption::ReuseMismatchedTypeRegisters}};
 
   auto options = parseEnvOptions("DISABLE", available_options);
 

csrc/options.h

@@ -109,6 +109,8 @@ enum class DisableOption {
                //! input shapes
   VarNameRemapping, //! Disable variable name remapping
   WelfordVectorization, //! Disable vectorizaton of Welford ops
+  ReuseMismatchedTypeRegisters, //! Disable explicitly re-using registers unless
+                                //! types match
   EndOfOption //! Placeholder for counting the number of elements
 };