Fix AMX support using MultiRamp

src/Deinterleave.cpp

@@ -610,7 +610,20 @@
         return expr;
     }
 
+    Scope<MemoryType> allocation_scope;
+    Stmt visit(const Allocate *op) override {
+        ScopedBinding<MemoryType> bind(allocation_scope, op->name, op->memory_type);
+        return IRMutator::visit(op);
+    }
+
     Stmt visit(const Store *op) override {
+        // Don't mess with matrix multiply ops, which use natively-supported 2D
+        // loads and stores.
+        if (auto *alloc = allocation_scope.find(op->name);
+            alloc && (*alloc) == MemoryType::AMXTile) {
+            return op;
+        }
+
         bool old_should_deinterleave = should_deinterleave;
         int old_num_lanes = num_lanes;
 
@@ -657,6 +670,13 @@
             return Stmt();
         }
 
+        // Don't mess with matrix multiply ops, which use natively-supported 2D
+        // loads and stores.
+        if (auto *alloc = allocation_scope.find(store->name);
+            alloc && (*alloc) == MemoryType::AMXTile) {
+            return Stmt();
+        }
+
         const Ramp *r0 = store->index.as<Ramp>();
 
         // It's not a store of a ramp index.