diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -352,8 +352,9 @@
     return *static_cast<fir::LLVMTypeConverter *>(this->getTypeConverter());
   }
 
-  void attachTBAATag(mlir::Operation *op, mlir::Type baseFIRType,
-                     mlir::Type accessFIRType, mlir::LLVM::GEPOp gep) const {
+  void attachTBAATag(mlir::LLVM::AliasAnalysisOpInterface op,
+                     mlir::Type baseFIRType, mlir::Type accessFIRType,
+                     mlir::LLVM::GEPOp gep) const {
     lowerTy().attachTBAATag(op, baseFIRType, accessFIRType, gep);
   }
 
diff --git a/flang/lib/Optimizer/CodeGen/TBAABuilder.h b/flang/lib/Optimizer/CodeGen/TBAABuilder.h
--- a/flang/lib/Optimizer/CodeGen/TBAABuilder.h
+++ b/flang/lib/Optimizer/CodeGen/TBAABuilder.h
@@ -170,8 +170,9 @@
 
   // Attach the llvm.tbaa attribute to the given memory accessing operation
   // based on the provided base/access FIR types and the GEPOp.
-  void attachTBAATag(mlir::Operation *op, mlir::Type baseFIRType,
-                     mlir::Type accessFIRType, mlir::LLVM::GEPOp gep);
+  void attachTBAATag(mlir::LLVM::AliasAnalysisOpInterface op,
+                     mlir::Type baseFIRType, mlir::Type accessFIRType,
+                     mlir::LLVM::GEPOp gep);
 
 private:
   // Return unique string name based on `basename`.
diff --git a/flang/lib/Optimizer/CodeGen/TBAABuilder.cpp b/flang/lib/Optimizer/CodeGen/TBAABuilder.cpp
--- a/flang/lib/Optimizer/CodeGen/TBAABuilder.cpp
+++ b/flang/lib/Optimizer/CodeGen/TBAABuilder.cpp
@@ -141,7 +141,7 @@
   return getAnyDataAccessTag();
 }
 
-void TBAABuilder::attachTBAATag(Operation *op, Type baseFIRType,
+void TBAABuilder::attachTBAATag(AliasAnalysisOpInterface op, Type baseFIRType,
                                 Type accessFIRType, GEPOp gep) {
   if (!enableTBAA)
     return;
@@ -163,10 +163,7 @@
   if (!tbaaTagSym)
     return;
 
-  auto tbaaAttr = ArrayAttr::get(op->getContext(), tbaaTagSym);
-  llvm::TypeSwitch<Operation *>(op)
-      .Case<LoadOp, StoreOp>([&](auto memOp) { memOp.setTbaaAttr(tbaaAttr); })
-      .Default([](auto) { llvm_unreachable("expected LoadOp or StoreOp"); });
+  op.setTBAATags(ArrayAttr::get(op->getContext(), tbaaTagSym));
 }
 
 } // namespace fir
diff --git a/flang/lib/Optimizer/CodeGen/TypeConverter.h b/flang/lib/Optimizer/CodeGen/TypeConverter.h
--- a/flang/lib/Optimizer/CodeGen/TypeConverter.h
+++ b/flang/lib/Optimizer/CodeGen/TypeConverter.h
@@ -390,8 +390,9 @@
   KindMapping &getKindMap() { return kindMapping; }
 
   // Relay TBAA tag attachment to TBAABuilder.
-  void attachTBAATag(mlir::Operation *op, mlir::Type baseFIRType,
-                     mlir::Type accessFIRType, mlir::LLVM::GEPOp gep) {
+  void attachTBAATag(mlir::LLVM::AliasAnalysisOpInterface op,
+                     mlir::Type baseFIRType, mlir::Type accessFIRType,
+                     mlir::LLVM::GEPOp gep) {
     tbaaBuilder.attachTBAATag(op, baseFIRType, accessFIRType, gep);
   }